Update app.py

app.py

@@ -2,17 +2,15 @@ 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_qdrant import QdrantVectorStore
+from langchain_chroma import Chroma
 import Raptor
-from io import StringIO
-from qdrant_client import QdrantClient
-from qdrant_client.models import Distance, VectorParams
 
 page = st.title("Chat with AskUSTH")
 
@@ -52,50 +50,6 @@ 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
-
-@st.cache_resource
-def update_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" 
-    )
-    
-    try:
-        client.delete_collection(collection_name=collection_name)
-    except Exception as e:
-        print(f"Warning: {e}")
-
-    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
 
@@ -123,12 +77,64 @@ 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 rag_chain(_model, _vectorstore):
-    retriever = _vectorstore.as_retriever()
+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()
     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
@@ -139,93 +145,24 @@ def rag_chain(_model, _vectorstore):
         {question}
         """
     prompt = PromptTemplate(template=template, input_variables=["context", "question"])
-    rag = (
+    rag_chain = (
         {"context": retriever | format_docs, "question": RunnablePassthrough()}
         | prompt
         | _model
         | StrOutputParser()
     )
-    return rag
-
-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 "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.add_texts(texts=all_texts)
-    rag = rag_chain(_model, _vectorstore)
-    return rag
-
-def reset_rag_chain(_model, _vectorstore):
-    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
+    return rag_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.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()  
-    
-@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
-    st.session_state.vector_store = update_chromadb("data")
-    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.uploaded_files and st.session_state.model is not None:
+    if st.session_state.rag is None:
+        load_rag()
+
 if st.session_state.model is not None:
     if st.session_state.llm is None:
         mess = ChatPromptTemplate.from_messages(
@@ -256,16 +193,12 @@ if st.session_state.model is not None:
             st.write(prompt)
             
         with st.chat_message("assistant"):
-            router = st.session_state.query_router.invoke(prompt)
-            switch = router.content
-            if "university" in switch:
+            if st.session_state.rag is not None:
                 respone = st.session_state.rag.invoke(prompt)
-                f_response = f"RAG: {respone}"
-                st.write(f_response)
+                st.write(respone)
             else: 
-                respone = st.session_state.llm.invoke(prompt)
-                f_response = f"other: {respone.content}"
-                st.write(f_response)
+                ans = st.session_state.llm.invoke(prompt)
+                respone = ans.content
+                st.write(respone)
                 
-        st.session_state.chat_history.append({"role": "assistant", "content": f_response})
-
+        st.session_state.chat_history.append({"role": "assistant", "content": respone})