Update app.py

app.py

@@ -10,6 +10,8 @@ from helpers import (
     list_docx_files, get_splits, get_json_splits_only, get_answer
 )
 
+from file_loader import get_vectorstore
+
 os.environ["USER_AGENT"] = "gradio-hf-space"
 
 if "GOOGLE_API_KEY" not in os.environ:
@@ -20,6 +22,70 @@ 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)
+
+    retriever = vectorstore.as_retriever(search_kwargs={"k": k})
+    results = retriever.invoke(query)
+
+    if results:
+        source_knowledge = "\n\n".join([doc.page_content for doc in results])
+        return f"""Using the contexts below, answer the query.
+Contexts:
+{source_knowledge}
+"""
+    else:
+        return f"No relevant context found.\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 energetic, helpful, polite, Vietnamese-speaking assistant. If you can not see the answer in contexts, try to search it up online by yourself but remember to give the source."}]},
+    {"role": "user", "parts": [{"text": augment_prompt(query)}]}
+]
+#     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
+#   '''
+
+    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
+    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)
+
 institutions = ['Tất cả'] + ['Trường Công Nghệ']
 categories = ['Tất cả'] + ['Đề án', 'Chương trình đào tạo']
 
@@ -39,4 +105,4 @@ with gr.Blocks() as demo:
                                       )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch()