Update app.py

app.py

@@ -13,6 +13,9 @@ from transformers import pipeline
 from langchain.llms import HuggingFaceHub
 from langchain.llms import HuggingFaceHub
 from langchain_core.retrievers import BaseRetriever
+from pydantic import BaseModel, Field
+from typing import List
+from typing import List, Dict, Any
 
 # Environment variables and configurations
 huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
@@ -33,10 +36,16 @@ def get_web_search_results(query: str, max_results: int = 10) -> List[Dict[str,
         print(f"An error occurred during web search: {str(e)}")
         return [{"error": f"An error occurred during web search: {str(e)}"}]
 
-class DuckDuckGoRetriever(BaseRetriever):
-    def __init__(self, embeddings):
-        self.embeddings = embeddings
-        self.vectorstore = FAISS.from_texts(["Initial document"], embeddings)
+class DuckDuckGoRetriever(BaseRetriever, BaseModel):
+    embeddings: Any = Field(description="The embeddings to use for the retriever")
+    vectorstore: Any = Field(default_factory=lambda: None, description="The vector store to use for the retriever")
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        self.vectorstore = FAISS.from_texts(["Initial document"], self.embeddings)
 
     def get_relevant_documents(self, query: str) -> List[Document]:
         search_results = get_web_search_results(query)
@@ -49,6 +58,9 @@ class DuckDuckGoRetriever(BaseRetriever):
         # Perform similarity search to get most relevant documents
         return self.vectorstore.similarity_search(query, k=3)
 
+    async def aget_relevant_documents(self, query: str) -> List[Document]:
+        return self.get_relevant_documents(query)
+
 def setup_retrieval_chain(model_name):
     # Set up the language model using HuggingFaceHub
     llm = HuggingFaceHub(
@@ -61,7 +73,7 @@ def setup_retrieval_chain(model_name):
     embeddings = HuggingFaceEmbeddings()
 
     # Create the DuckDuckGo retriever
-    retriever = DuckDuckGoRetriever(embeddings)
+    retriever = DuckDuckGoRetriever(embeddings=embeddings)
 
     # Set up the memory
     memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)