Update app.py

app.py

@@ -15,24 +15,19 @@ from helpers import (
 from file_loader import get_vectorstore
 # import asyncio
 
-if "GOOGLE_API_KEY" not in os.environ:
-    os.environ["GOOGLE_API_KEY"] = "AIzaSyDJ4vIKuIBIPNHATLxnoHlagXWbsAz-vRs"
 key = "AIzaSyDJ4vIKuIBIPNHATLxnoHlagXWbsAz-vRs"
-
-###
-
 # Cấu hình API key cho Google GenAI
 genai.configure(api_key=key)
 
 vectorstore = get_vectorstore()
-# Define the augment_prompt function
-def augment_prompt(query: str, k: int = 10):
-    queries = []
-    queries.append(query)
+last_vector_docs = None  # Lưu kết quả docs từ vectorstore.invoke trong lần gọi get_answer gần nhất
 
+def augment_prompt(query: str, k: int = 10):
+    global last_vector_docs
     retriever = vectorstore.as_retriever(search_kwargs={"k": k})
     results = retriever.invoke(query)
-
+    # Lưu kết quả để dùng cho log và lọc sau này
+    last_vector_docs = results  
     if results:
         source_knowledge = "\n\n".join([doc.page_content for doc in results])
         return f"""Dữ liệu dưới đây liên quan đến Trường Công Nghệ (NCT) thuộc Đại học Kinh tế Quốc dân (NEU), dựa vào đó trả lời câu hỏi.
@@ -40,53 +35,49 @@ Dữ liệu:
 {source_knowledge}
 """
     else:
-        return f"Không có thông tin liên quan. \n."
+        return "Không có thông tin liên quan.\n."
 
 def get_answer(query, queries_list=None):
     if queries_list is None:
         queries_list = []
     
     messages = [
-    {"role": "user", "parts": [{"text": "IMPORTANT: You are a super helpful, polite, Vietnamese-speaking asistant to give infomation of an university. If you can not see the answer in contexts, try to search it up online by yourself OR tell user to make a more detailed question."}]},
-    {"role": "user", "parts": [{"text": augment_prompt(query=query,k=100)}]}
-]
-#     bonus = '''
-# Bạn tham kháo thêm các nguồn thông tin tại:
-# Trang thông tin điện tử: https://neu.edu.vn ; https://daotao.neu.edu.vn
-# Trang mạng xã hội có thông tin tuyển sinh: https://www.facebook.com/ktqdNEU ; https://www.facebook.com/tvtsneu ;
-# Email tuyển sinh: [email protected]
-# Số điện thoại tuyển sinh: 0888.128.558
-#   '''
-
+        {"role": "user", "parts": [{"text": "IMPORTANT: You are a super helpful, polite, Vietnamese-speaking assistant to give information of an university. If you cannot see the answer in contexts, try to search it up online by yourself OR tell user to make a more detailed question."}]},
+        {"role": "user", "parts": [{"text": augment_prompt(query=query, k=20)}]}
+    ]
+    
     queries_list.append(query)
     queries = {"role": "user", "parts": [{"text": prompt_order(queries_list)}]}
     messages_with_queries = messages.copy()
     messages_with_queries.append(queries)
-    # messages_with_queries.insert(0, queries)
-
-  # Configure API key
+    
+    # Cấu hình API key và khởi tạo model Gemini
     genai.configure(api_key=key)
-
-  # Initialize the Gemini model
     model = genai.GenerativeModel("gemini-2.0-flash")
-
+    
     response = model.generate_content(contents=messages_with_queries, stream=True)
     response_text = ""
-
     for chunk in response:
         response_text += chunk.text
         yield response_text
-
     messages.append({"role": "model", "parts": [{"text": response_text}]})
-
-        # user_feedback = yield "\nNhập phản hồi của bạn (hoặc nhập 'q' để thoát): "
-        # if user_feedback.lower() == "q":
-        #     break
-
-        # messages.append({"role": "user", "parts": [{"text": query}]})
-
     log_message(messages)
 
+def filter_vector_docs(keyword: str):
+    global last_vector_docs
+    if last_vector_docs is None:
+        return "Chưa có dữ liệu vectorstore được gọi từ get_answer."
+    else:
+        if not keyword.strip():
+            # Nếu không nhập gì, trả về tất cả
+            filtered = [doc.page_content for doc in last_vector_docs]
+        else:
+            # Lọc các chunk chứa từ khoá (không phân biệt chữ hoa thường)
+            filtered = [doc.page_content for doc in last_vector_docs if keyword.lower() in doc.page_content.lower()]
+        if not filtered:
+            return f"Không có kết quả chứa từ khoá '{keyword}'."
+        return "\n\n".join(filtered)
+
 institutions = ['Tất cả', 'Trường Công Nghệ']
 categories = ['Tất cả', 'Đề án', 'Chương trình đào tạo']
 
@@ -94,20 +85,21 @@ print("Launching on space... This may take some time...")
 
 with gr.Blocks() as demo:
     with gr.Row():
-        category1 = gr.Dropdown(choices = institutions, label="Trường", value = None)
-        category2 = gr.Dropdown(choices = categories, label="Bạn quan tâm tới", value = None)
-
-    chat_interface = gr.ChatInterface(get_answer,
-                                      textbox=gr.Textbox(placeholder="Đặt câu hỏi tại đây",
-                                                        container=False,
-                                                        autoscroll=True,
-                                                        scale=7),
-                                      type="messages",
-                                      # textbox=prompt,
-                                      # additional_inputs=[category1, category2]
-                                      )
-# playwright_path = get_driver_dir()
+        # Dropdown category nếu cần
+        category1 = gr.Dropdown(choices=institutions, label="Trường", value=None)
+        category2 = gr.Dropdown(choices=categories, label="Bạn quan tâm tới", value=None)
+    
+    # Chat Interface sử dụng ô nhập chung
+    shared_query = gr.Textbox(placeholder="Đặt câu hỏi tại đây", container=False, autoscroll=True, scale=7)
+    chat_interface = gr.ChatInterface(get_answer, textbox=shared_query, type="messages")
+    
+    # Phần lọc các chunk: ô prompt nhập từ khoá và nút "Tìm trích xuất" nằm cùng hàng,
+    # kết quả sẽ hiển thị ở ô bên dưới. Nếu để trống, hiển thị toàn bộ.
+    with gr.Row():
+        filter_prompt = gr.Textbox(label="Nhập từ khoá", placeholder="Nhập từ khoá", interactive=True)
+        filter_button = gr.Button("Đọc trích xuất")
+    filter_output = gr.Textbox(label="Content", interactive=False)
+    filter_button.click(fn=filter_vector_docs, inputs=filter_prompt, outputs=filter_output)
 
 if __name__ == "__main__":
-    demo.launch()    
-    # demo.launch()
+    demo.launch()