Update chatbot/core.py

chatbot/core.py

@@ -3,7 +3,9 @@ from chatbot.memory import memory
 from chatbot.prompts import chat_prompt
 from langchain.retrievers import WikipediaRetriever
 from langchain.chains import ConversationalRetrievalChain
-from langchain.schema import BaseRetriever
+from pydantic import Field
+from typing import List, Callable
+from langchain.schema import BaseRetriever, Document
 
 def translate_to_english(text: str) -> str:
     """Use Gemini LLM to translate text to English."""
@@ -12,26 +14,25 @@ def translate_to_english(text: str) -> str:
     return response  # Assuming `gemini_llm.invoke()` returns plain text
 
 class WikipediaTranslationRetriever(BaseRetriever):
-    """Custom Retriever that translates queries before searching Wikipedia."""
-    def __init__(self, retriever, translator):
-        self.retriever = retriever
-        self.translator = translator
+    retriever: WikipediaRetriever = Field(..., description="The underlying Wikipedia retriever")
+    translator: Callable[[str], str] = Field(..., description="Function to translate queries to English")
 
-    def get_relevant_documents(self, query):
-        translated_query = self.translator(query)  # Translate query to English
+    def get_relevant_documents(self, query: str) -> List[Document]:
+        translated_query = self.translator(query)
         print(f"🔄 Translated Query: {translated_query}")
         return self.retriever.get_relevant_documents(translated_query)
+    
+    async def aget_relevant_documents(self, query: str) -> List[Document]:
+        # For simplicity, we are not implementing the async version.
+        raise NotImplementedError("Async retrieval is not implemented.")
 
-    async def aget_relevant_documents(self, query):
-        # If your environment doesn't need async support, you can simply raise an error.
-        raise NotImplementedError("Asynchronous retrieval is not implemented.")
 
+# Create the retriever instance to be used in your qa_chain:
+retriever = WikipediaTranslationRetriever(
+    retriever=WikipediaRetriever(),
+    translator=translate_to_english
+)
 
-# ✅ Use WikipediaRetriever
-wiki_retriever = WikipediaRetriever()
-
-# ✅ Wrap with translation
-retriever = WikipediaTranslationRetriever(wiki_retriever, translate_to_english)
 
 # ✅ Use ConversationalRetrievalChain
 qa_chain = ConversationalRetrievalChain.from_llm(