Update app.py

app.py

@@ -10,17 +10,32 @@ from langchain_core.runnables import RunnablePassthrough
 from langchain_chroma import Chroma
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 
+# App Title
 page = st.title("Chat with AskUSTH")
 
+# Initialize session states
 if "gemini_api" not in st.session_state:
     st.session_state.gemini_api = None
 
 if "rag" not in st.session_state:
     st.session_state.rag = None
-    
+
 if "llm" not in st.session_state:
     st.session_state.llm = None
 
+if "embd" not in st.session_state:
+    st.session_state.embd = None
+
+if "model" not in st.session_state:
+    st.session_state.model = None
+
+if "save_dir" not in st.session_state:
+    st.session_state.save_dir = None 
+
+if "uploaded_files" not in st.session_state:
+    st.session_state.uploaded_files = set()
+
+# Caching functions
 @st.cache_resource
 def get_chat_google_model(api_key):
     os.environ["GOOGLE_API_KEY"] = api_key
@@ -42,23 +57,56 @@ def get_embedding_model():
         model_name=model_name,
         model_kwargs=model_kwargs,
         encode_kwargs=encode_kwargs
-    )   
+    )
     return model
 
-if "embd" not in st.session_state:
-    st.session_state.embd = get_embedding_model()
+# Load and process text files
+def load_txt(file_path):
+    loader = TextLoader(file_path=file_path, encoding="utf-8")
+    doc = loader.load()
+    return doc
 
-if "model" not in st.session_state:
-    st.session_state.model = None
+def format_docs(docs):
+    return "\n\n".join(doc.page_content for doc in docs)
 
-if "save_dir" not in st.session_state:
-    st.session_state.save_dir = None 
+# Compute RAG Chain
+@st.cache_resource
+def compute_rag_chain(_model, _embd, docs_texts):
+    if not docs_texts:
+        raise ValueError("No documents to process. Please upload valid text files.")
+    
+    combined_text = "\n\n".join(docs_texts)
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10)
+    texts = text_splitter.split_text(combined_text)
+    
+    if not texts:
+        raise ValueError("Text splitter did not generate any text chunks. Check your input.")
+    
+    vectorstore = Chroma.from_texts(texts=texts, embedding=_embd)
+    retriever = vectorstore.as_retriever()
 
-if "uploaded_files" not in st.session_state:
-    st.session_state.uploaded_files = set()
+    template = """
+        Bạn là một trợ lí AI hỗ trợ tuyển sinh và sinh viê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ếu bạn không biết câu trả lời, hãy nói không biết, đừng cố tạo ra câu trả lời.
+        Dưới đây là thông tin liên quan mà bạn cần sử dụng tới:
+        {context}
+        hãy trả lời:
+        {question}
+    """
+    prompt = PromptTemplate(template=template, input_variables=["context", "question"])
     
+    rag_chain = (
+        {"context": retriever | format_docs, "question": RunnablePassthrough()}
+        | prompt
+        | _model
+        | StrOutputParser()
+    )
+    return rag_chain
+
+# Dialog to setup Gemini
 @st.dialog("Setup Gemini")
-def vote():
+def setup_gemini():
     st.markdown(
         """
         Để sử dụng Google Gemini, bạn cần cung cấp API key. Tạo key của bạn [tại đây](https://ai.google.dev/gemini-api/docs/get-started/tutorial?lang=python&hl=vi) và dán vào bên dưới.
@@ -67,115 +115,44 @@ def vote():
     key = st.text_input("Key:", "")
     if st.button("Save") and key != "":
         st.session_state.gemini_api = key
-        st.rerun()  
+        st.rerun()
 
 if st.session_state.gemini_api is None:
-    vote()
+    setup_gemini()
 
 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.embd is None:
+    st.session_state.embd = get_embedding_model()
+
 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
 
+# Sidebar to upload files
 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
-                
+    uploaded_files = st.file_uploader("Chọn file txt", accept_multiple_files=True, type=["txt"])
+    if uploaded_files:
+        documents = []
+        uploaded_file_names = set()
+        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())
                 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):
-    # Combine all texts into one large string
-    combined_text = "\n\n".join(docs_texts)  # Join all document texts into one string
-
-    # Use RecursiveCharacterTextSplitter to split text into chunks
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=0)
-    texts = text_splitter.split_text(combined_text)  # Now this will work as 'combined_text' is a string
-
-    # Create vector store for similarity search
-    vectorstore = Chroma.from_texts(texts=texts, embedding=_embd)
-    retriever = vectorstore.as_retriever()
-
-    # Prepare the prompt for context and question
-    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
-        Nếu bạn không biết câu trả lời, hãy nói không biết, đừng cố tạo ra câu trả lời.\n
-        Dưới đây là thông tin liên quan mà bạn cần sử dụng tới:\n
-        {context}\n
-        hãy trả lời:\n
-        {question}
-        """
-    prompt = PromptTemplate(template=template, input_variables=["context", "question"])
-
-    # Chain for RAG
-    rag_chain = (
-        {"context": retriever | format_docs, "question": RunnablePassthrough()}
-        | prompt
-        | _model
-        | StrOutputParser()
-    )
-    return rag_chain
-
-@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.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(
-            [
-                (
-                    "system",
-                    "Bản là một trợ lí AI hỗ trợ tuyển sinh và sinh viên",
-                ),
-                ("human", "{input}"),
-            ]
-        )
-        chain = mess | st.session_state.model
-        st.session_state.llm = chain
+        
+        if documents:
+            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.session_state.uploaded_files = uploaded_file_names
 
+# Chat Interface
 if "chat_history" not in st.session_state:
     st.session_state.chat_history = []
 
@@ -184,20 +161,14 @@ for message in st.session_state.chat_history:
         st.write(message["content"])
 
 prompt = st.chat_input("Bạn muốn hỏi gì?")
-if st.session_state.model is not None:
-    if prompt:
-        st.session_state.chat_history.append({"role": "user", "content": prompt})
-            
-        with st.chat_message("user"):
-            st.write(prompt)
-            
-        with st.chat_message("assistant"):
-            if st.session_state.rag is not None:
-                respone = st.session_state.rag.invoke(prompt)
-                st.write(respone)
-            else: 
-                ans = st.session_state.llm.invoke(prompt)
-                respone = ans.content
-                st.write(respone)
-                
-        st.session_state.chat_history.append({"role": "assistant", "content": respone})
+if prompt and st.session_state.model:
+    st.session_state.chat_history.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.write(prompt)
+    with st.chat_message("assistant"):
+        if st.session_state.rag:
+            response = st.session_state.rag.invoke(prompt)
+        else:
+            response = st.session_state.model.invoke(prompt).content
+        st.write(response)
+        st.session_state.chat_history.append({"role": "assistant", "content": response})