v1

app.py

@@ -2,15 +2,17 @@ import streamlit as st
 import os
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_community.document_loaders import TextLoader
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain.prompts import PromptTemplate
 
 from langchain_core.output_parsers import StrOutputParser
 
 from langchain_core.runnables import RunnablePassthrough
-from langchain_chroma import Chroma
+from langchain_qdrant import QdrantVectorStore
 import Raptor
+from io import StringIO
+from qdrant_client import QdrantClient
+from qdrant_client.models import Distance, VectorParams
 
 page = st.title("Chat with AskUSTH")
 
@@ -23,6 +25,8 @@ if "rag" not in st.session_state:
 if "llm" not in st.session_state:
     st.session_state.llm = None
 
+
+
 @st.cache_resource
 def get_chat_google_model(api_key):
     os.environ["GOOGLE_API_KEY"] = api_key
@@ -50,6 +54,27 @@ def get_embedding_model():
 if "embd" not in st.session_state:
     st.session_state.embd = get_embedding_model()
 
+@st.cache_resource
+def load_chromadb(collection_name):
+    client = QdrantClient(
+    url="https://da9fadd2-dc5a-4481-ac79-4e2677a2354b.europe-west3-0.gcp.cloud.qdrant.io", 
+    api_key="X_-IVToBM07Mot4Mmzg5xNjYzc1DlIgl0VQDUNmGhI_Z-WA5FJ2ETA" 
+)
+
+    client.recreate_collection(
+        collection_name=collection_name, 
+        vectors_config=VectorParams(size=768, distance=Distance.COSINE) 
+    )
+    db = QdrantVectorStore(
+        client=client,
+        collection_name=collection_name,
+        embedding=st.session_state.embd,
+    )
+    return db
+
+if "vector_store" not in st.session_state:
+    st.session_state.vector_store = load_chromadb("data")
+
 if "model" not in st.session_state:
     st.session_state.model = None
 
@@ -77,64 +102,12 @@ if st.session_state.gemini_api is None:
 if st.session_state.gemini_api and st.session_state.model is None:
     st.session_state.model = get_chat_google_model(st.session_state.gemini_api)
 
-if st.session_state.save_dir is None:
-    save_dir = "./Documents"
-    if not os.path.exists(save_dir):
-        os.makedirs(save_dir)
-    st.session_state.save_dir = save_dir
-    
-def load_txt(file_path):
-    loader_sv = TextLoader(file_path=file_path, encoding="utf-8")
-    doc = loader_sv.load()
-    return doc
-
-with st.sidebar:
-    uploaded_files = st.file_uploader("Chọn file txt", accept_multiple_files=True, type=["txt"])   
-    if st.session_state.gemini_api:
-        if uploaded_files:
-            documents = []
-            uploaded_file_names = set()
-            new_docs = False
-            for uploaded_file in uploaded_files:
-                uploaded_file_names.add(uploaded_file.name)
-                if uploaded_file.name not in st.session_state.uploaded_files:
-                    file_path = os.path.join(st.session_state.save_dir, uploaded_file.name)
-                    with open(file_path, mode='wb') as w:
-                        w.write(uploaded_file.getvalue())
-                else:
-                    continue
-
-                new_docs = True
-                
-                doc = load_txt(file_path)
-                
-                documents.extend([*doc])
-                
-            if new_docs:  
-                st.session_state.uploaded_files = uploaded_file_names
-                st.session_state.rag = None
-        else:
-            st.session_state.uploaded_files = set()
-            st.session_state.rag = None
-
 def format_docs(docs):
     return "\n\n".join(doc.page_content for doc in docs)
 
 @st.cache_resource
-def compute_rag_chain(_model, _embd, docs_texts):
-    results = Raptor.recursive_embed_cluster_summarize(_model, _embd, docs_texts, level=1, n_levels=3)
-    all_texts = docs_texts.copy()
-    i = 0
-    for level in sorted(results.keys()):
-        summaries = results[level][1]["summaries"].tolist()
-        all_texts.extend(summaries)
-        print(f"summary {i} -------------------------------------------------")
-        print(summaries)
-        i += 1
-    print("all_texts ______________________________________")
-    print(all_texts)
-    vectorstore = Chroma.from_texts(texts=all_texts, embedding=_embd)
-    retriever = vectorstore.as_retriever()
+def rag_chain(_model, _vectorstore):
+    retriever = _vectorstore.as_retriever()
     template = """
         Bạn là một trợ lí AI hỗ trợ tuyển sinh và sinh viên. \n
         Hãy trả lời câu hỏi chính xác, tập trung vào thông tin liên quan đến câu hỏi. \n
@@ -145,24 +118,105 @@ def compute_rag_chain(_model, _embd, docs_texts):
         {question}
         """
     prompt = PromptTemplate(template=template, input_variables=["context", "question"])
-    rag_chain = (
+    rag = (
         {"context": retriever | format_docs, "question": RunnablePassthrough()}
         | prompt
         | _model
         | StrOutputParser()
     )
-    return rag_chain
+    return rag
 
-@st.dialog("Setup RAG")
-def load_rag():
-    docs_texts = [d.page_content for d in documents]
-    st.session_state.rag = compute_rag_chain(st.session_state.model, st.session_state.embd, docs_texts)
-    st.rerun()  
+if st.session_state.model is not None and st.session_state.vector_store is not None:
+    st.session_state.rag = rag_chain(st.session_state.model, st.session_state.vector_store)
 
-if st.session_state.uploaded_files and st.session_state.model is not None:
-    if st.session_state.rag is None:
-        load_rag()
+# if st.session_state.save_dir is None:
+#     save_dir = "./Documents"
+#     if not os.path.exists(save_dir):
+#         os.makedirs(save_dir)
+#     st.session_state.save_dir = save_dir
+    
+# def load_txt(file_path):
+#     loader_sv = TextLoader(file_path=file_path, encoding="utf-8")
+#     doc = loader_sv.load()
+#     return doc
+
+if "new_docs" not in st.session_state:
+    st.session_state.new_docs = False
+
+with st.sidebar:
+    uploaded_files = st.file_uploader("Chọn file txt", accept_multiple_files=True, type=["txt"])   
+    if st.session_state.model:
+        documents = []
+        uploaded_file_names = set()
+        if uploaded_files:
+            for uploaded_file in uploaded_files:
+                uploaded_file_names.add(uploaded_file.name)      
+        if uploaded_file_names != st.session_state.uploaded_files and not st.session_state.new_docs:
+            st.session_state.uploaded_files = uploaded_file_names
+            st.session_state.new_docs = True
+            if uploaded_files:
+                for uploaded_file in uploaded_files: 
+                    stringio=StringIO(uploaded_file.getvalue().decode('utf-8'))
+                    read_data=str(stringio.read())
+                    documents.append(read_data)
+                
+def update_rag_chain(_model, _embd, _vectorstore, docs_texts):
+    results = Raptor.recursive_embed_cluster_summarize(_model, _embd, docs_texts, level=1, n_levels=3)
+    all_texts = docs_texts.copy()
+    for level in sorted(results.keys()):
+        summaries = results[level][1]["summaries"].tolist()
+        all_texts.extend(summaries)
+    _vectorstore.reset_collection()
+    _vectorstore.add_texts(texts=all_texts)
+    rag = rag_chain(_model, _vectorstore)
+    return rag
+
+def reset_rag_chain(_model, _vectorstore):
+    _vectorstore.reset_collection()
+    rag = rag_chain(_model, _vectorstore)
+    return rag
+
+if "query_router" not in st.session_state:
+    st.session_state.query_router = None  
+
+@st.cache_resource
+def query_router(_model):
+    mess = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                """Bạn là một chatbot hỗ trợ giải đáp về đại học, nhiệm vụ của bạn là phân loại câu hỏi. 
+                Nếu câu hỏi về đại học thì trả về 'university', nếu không liên quan tới tuyển sinh và sinh viên thì trả về 'other'. 
+                Bắt buộc Kết quả chỉ trả về với một trong hai lựa chọn trên. 
+                Không được trả lời thêm bất kỳ thông tin nào khác.""",
+            ),
+            ("human", "{input}"),
+        ]
+    )
+    chain = mess | _model
+    return chain
 
+if st.session_state.model is not None:
+    st.session_state.query_router = query_router(st.session_state.model)
+
+@st.dialog("Update DB")
+def update_vectorstore(_model, _embd, _vectorstore, docs):
+    docs_texts = [d for d in docs]
+    st.session_state.rag = update_rag_chain(_model, _embd, _vectorstore, docs_texts)   
+    st.rerun()  
+    
+@st.dialog("Reset DB")
+def reset_vectorstore(_model, _vectorstore):
+    st.session_state.rag = reset_rag_chain(_model, _vectorstore) 
+    st.rerun()
+
+if st.session_state.new_docs:
+    st.session_state.new_docs = False
+    if st.session_state.uploaded_files:
+        update_vectorstore(st.session_state.model, st.session_state.embd, st.session_state.vector_store, documents)
+    else:
+        reset_vectorstore(st.session_state.model, st.session_state.vector_store)
+    
 if st.session_state.model is not None:
     if st.session_state.llm is None:
         mess = ChatPromptTemplate.from_messages(
@@ -193,13 +247,16 @@ if st.session_state.model is not None:
             st.write(prompt)
             
         with st.chat_message("assistant"):
-            if st.session_state.rag is not None:
+            router = st.session_state.query_router.invoke(prompt)
+            switch = router.content
+            if "university" in switch:
                 respone = st.session_state.rag.invoke(prompt)
-                st.write(respone)
+                f_response = f"RAG: {respone}"
+                st.write(f_response)
             else: 
-                ans = st.session_state.llm.invoke(prompt)
-                respone = ans.content
-                st.write(respone)
+                respone = st.session_state.llm.invoke(prompt)
+                f_response = f"other: {respone.content}"
+                st.write(f_response)
                 
-        st.session_state.chat_history.append({"role": "assistant", "content": respone})
+        st.session_state.chat_history.append({"role": "assistant", "content": f_response})
 

requirements.txt

@@ -4,4 +4,5 @@ langchain-community
 langchain-huggingface
 umap-learn
 scikit-learn
-langchain-chroma
+langchain-qdrant
+qdrant-client