Update chatbot/core.py

chatbot/core.py

@@ -1,15 +1,36 @@
+from chatbot.llm import gemini_llm  # Import Gemini LLM
+from chatbot.memory import memory
+from chatbot.prompts import chat_prompt
 from langchain.retrievers import WikipediaRetriever
-from chatbot.llm import gemini_llm  
-from chatbot.memory import memory  
-from chatbot.prompts import chat_prompt  
 from langchain.chains import ConversationalRetrievalChain
 
-# Create the retriever
-retriever = WikipediaRetriever()
+def translate_to_english(text: str) -> str:
+    """Use Gemini LLM to translate text to English."""
+    prompt = f"Translate the following text to English:\n\n{text}"
+    response = gemini_llm.invoke(prompt)  # Invoke Gemini for translation
+    return response  # Assuming `gemini_llm.invoke()` returns plain text
 
+class WikipediaTranslationRetriever:
+    """Custom Retriever that translates queries before searching Wikipedia."""
+    def __init__(self, retriever, translator):
+        self.retriever = retriever
+        self.translator = translator
+
+    def get_relevant_documents(self, query):
+        translated_query = self.translator(query)  # Translate query to English
+        print(f"🔄 Translated Query: {translated_query}")
+        return self.retriever.get_relevant_documents(translated_query)
+
+# ✅ Use WikipediaRetriever
+wiki_retriever = WikipediaRetriever()
+
+# ✅ Wrap with translation
+retriever = WikipediaTranslationRetriever(wiki_retriever, translate_to_english)
+
+# ✅ Use ConversationalRetrievalChain
 qa_chain = ConversationalRetrievalChain.from_llm(
     llm=gemini_llm, 
-    retriever=retriever,  # Dynamic Wikipedia search
+    retriever=retriever, 
     memory=memory,  
     return_source_documents=False,  
     combine_docs_chain_kwargs={"prompt": chat_prompt},
@@ -17,10 +38,10 @@ qa_chain = ConversationalRetrievalChain.from_llm(
 )
 
 def get_chat_response(user_input: str) -> str:
-    response = qa_chain(user_input)
+    """Process user input and return chat response using Wikipedia retrieval."""
+    response = qa_chain(user_input)  # Pass query to retrieval-based QA chain
 
-    # Lưu vào bộ nhớ hội thoại
+    # Save conversation context
     memory.save_context({"input": user_input}, {"output": response["result"]})
 
     return response["result"]
-