Update app.py

app.py

@@ -10,10 +10,10 @@ from langchain_core.runnables import RunnablePassthrough
 from langchain_chroma import Chroma
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 
-# App Title
+# Tiêu đề ứng dụng
 page = st.title("Chat with AskUSTH")
 
-# Initialize session states
+# Khởi tạo trạng thái phiên
 if "gemini_api" not in st.session_state:
     st.session_state.gemini_api = None
 
@@ -35,6 +35,21 @@ if "save_dir" not in st.session_state:
 if "uploaded_files" not in st.session_state:
     st.session_state.uploaded_files = set()
 
+if "chat_history" not in st.session_state:
+    st.session_state.chat_history = []
+
+# Hàm tải và xử lý file văn bản
+def load_txt(file_path):
+    loader = TextLoader(file_path=file_path, encoding="utf-8")
+    doc = loader.load()
+    return doc
+
+# Hàm định dạng văn bản
+def format_docs(docs):
+    """Định dạng các tài liệu thành chuỗi văn bản."""
+    return "\n\n".join(doc.page_content for doc in docs)
+
+# Hàm thiết lập mô hình Google Gemini
 @st.cache_resource
 def get_chat_google_model(api_key):
     os.environ["GOOGLE_API_KEY"] = api_key
@@ -46,6 +61,7 @@ def get_chat_google_model(api_key):
         max_retries=2,
     )
 
+# Hàm thiết lập mô hình embedding
 @st.cache_resource
 def get_embedding_model():
     model_name = "bkai-foundation-models/vietnamese-bi-encoder"
@@ -59,26 +75,18 @@ def get_embedding_model():
     )
     return model
 
-def load_txt(file_path):
-    loader = TextLoader(file_path=file_path, encoding="utf-8")
-    doc = loader.load()
-    return doc
-
-def format_docs(docs):
-    """Format documents into a single string for prompt input."""
-    return "\n\n".join(doc.page_content for doc in docs)
-
+# Hàm tạo 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.")
+        raise ValueError("Không có tài liệu nào để xử lý. Vui lòng tải lên các tệp hợp lệ.")
     
     combined_text = "\n\n".join(docs_texts)
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
     texts = text_splitter.split_text(combined_text)
     
     if len(texts) > 5000:
-        raise ValueError("The document creates too many chunks. Please use smaller documents.")
+        raise ValueError("Tài liệu tạo ra quá nhiều đoạn. Vui lòng sử dụng tài liệu nhỏ hơn.")
     
     vectorstore = Chroma.from_texts(texts=texts, embedding=_embd)
     retriever = vectorstore.as_retriever()
@@ -102,7 +110,7 @@ def compute_rag_chain(_model, _embd, docs_texts):
     )
     return rag_chain
 
-# Dialog to setup Gemini
+# Dialog cài đặt Google Gemini
 @st.dialog("Setup Gemini")
 def setup_gemini():
     st.markdown(
@@ -130,6 +138,7 @@ if st.session_state.save_dir is None:
         os.makedirs(save_dir)
     st.session_state.save_dir = save_dir
 
+# Sidebar: Upload file và xử lý
 with st.sidebar:
     uploaded_files = st.file_uploader("Chọn file txt", accept_multiple_files=True, type=["txt"])
     max_file_size_mb = 5
@@ -150,3 +159,26 @@ with st.sidebar:
         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)
+
+# Giao diện chat
+for message in st.session_state.chat_history:
+    with st.chat_message(message["role"]):
+        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:
+                response = st.session_state.rag.invoke(prompt)
+                st.write(response)
+            else: 
+                ans = st.session_state.llm.invoke(prompt)
+                response = ans.content
+                st.write(response)
+                
+        st.session_state.chat_history.append({"role": "assistant", "content": response})