Use wikipedia both en and vi

chatbot/core.py

@@ -1,20 +1,30 @@
+from langchain.retrievers import WikipediaRetriever
+import wikipedia
 from chatbot.llm import gemini_llm  
-from chatbot.retrieval import get_vector_db
-from chatbot.memory import memory
-from chatbot.prompts import chat_prompt
+from chatbot.memory import memory  
+from chatbot.prompts import chat_prompt  
 from langchain.chains import ConversationalRetrievalChain
 
-vector_db = get_vector_db()
+def search_wikipedia(query: str, language: str = "vi"):
+    """Search Wikipedia in the specified language (vi or en)."""
+    wikipedia.set_lang(language)  # Set Wikipedia language dynamically
+    retriever = WikipediaRetriever()  # Create a new retriever each time to apply language setting
+    return retriever.get_relevant_documents(query)
 
-retriever = vector_db.as_retriever(search_kwargs={"k": 20})
+def get_retriever(user_input: str):
+    """Decide which language retriever to use based on user input."""
+    # Example logic: If input contains English words, use "en"; otherwise, use "vi".
+    if any(char.isascii() for char in user_input):
+        return search_wikipedia(user_input, language="en")
+    return search_wikipedia(user_input, language="vi")
 
 qa_chain = ConversationalRetrievalChain.from_llm(
     llm=gemini_llm, 
-    retriever=retriever, 
+    retriever=get_retriever,  # Dynamic Wikipedia search
     memory=memory,  
-    return_source_documents= False,  
+    return_source_documents=False,  
     combine_docs_chain_kwargs={"prompt": chat_prompt},
-    output_key="result" 
+    output_key="result"
 )
 
 def get_chat_response(user_input: str) -> str:
@@ -24,3 +34,4 @@ def get_chat_response(user_input: str) -> str:
     memory.save_context({"input": user_input}, {"output": response["result"]})
 
     return response["result"]
+