Update retriever.py

retriever.py

@@ -2,12 +2,14 @@ from smolagents import Tool
 from langchain_community.vectorstores import FAISS
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain.docstore.document import Document
+from langchain.retrievers import EnsembleRetriever
+from langchain_community.retrievers import BM25Retriever
 from tools import DuckDuckGoSearchTool
 import datasets
 
-class GuestInfoRetrieverTool(Tool):
-    name = "guest_info_retriever"
-    description = "Retrieves detailed information about gala guests using semantic search."
+class MultiIndexRetrieverTool(Tool):
+    name = "multi_index_guest_retriever"
+    description = "Retrieves guest information from multiple indexes and verified sources."
     inputs = {
         "query": {
             "type": "string",
@@ -16,27 +18,61 @@ class GuestInfoRetrieverTool(Tool):
     }
     output_type = "string"
 
-    def __init__(self, docs):
+    def __init__(self, primary_docs, secondary_docs=None):
         self.is_initialized = False
-        # Initialize embedding model
         self.embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-        # Create FAISS vector store
-        self.retriever = FAISS.from_documents(docs, self.embeddings).as_retriever(
-            search_kwargs={"k": 3})  # Return top 3 results
+        
+        # Primary index (guest dataset)
+        self.primary_retriever = FAISS.from_documents(
+            primary_docs, self.embeddings
+        ).as_retriever(search_kwargs={"k": 3})
+        
+        # Secondary index (e.g., Wikipedia or another dataset)
+        self.secondary_retriever = None
+        if secondary_docs:
+            self.secondary_retriever = FAISS.from_documents(
+                secondary_docs, self.embeddings
+            ).as_retriever(search_kwargs={"k": 3})
+        
+        # BM25 for keyword-based fallback
+        self.bm25_retriever = BM25Retriever.from_documents(primary_docs)
+        self.bm25_retriever.k = 3
+        
+        # Ensemble retriever (combines primary and secondary)
+        retrievers = [self.primary_retriever, self.bm25_retriever]
+        if self.secondary_retriever:
+            retrievers.append(self.secondary_retriever)
+        
+        self.ensemble_retriever = EnsembleRetriever(
+            retrievers=retrievers, weights=[0.5, 0.3, 0.2] if self.secondary_retriever else [0.7, 0.3]
+        )
+        
         self.web_search_tool = DuckDuckGoSearchTool()
-    
+
     def forward(self, query: str):
-        results = self.retriever.get_relevant_documents(query)
+        # Retrieve from ensemble
+        results = self.ensemble_retriever.get_relevant_documents(query)
+        
         if results:
-            return "\n\n".join([doc.page_content for doc in results])
-        else:
-            # Fallback to web search
-            web_results = self.web_search_tool.forward(f"Who is {query}?")
-            return f"No guest found in dataset. Web search results:\n{web_results}"
+            # Filter for verified sources (e.g., prioritize dataset over secondary)
+            verified_results = [
+                doc for doc in results if doc.metadata.get("source", "").startswith("unit3-invitees")
+            ]
+            other_results = [
+                doc for doc in results if not doc.metadata.get("source", "").startswith("unit3-invitees")
+            ]
+            combined_results = verified_results[:2] + other_results[:1]  # Prioritize verified
+            if combined_results:
+                return "\n\n".join([doc.page_content for doc in combined_results])
+        
+        # Fallback to web search
+        web_results = self.web_search_tool.forward(f"Who is {query}?")
+        return f"No guest found in indexes. Web search results:\n{web_results}"
 
 def load_guest_dataset():
+    # Primary dataset
     guest_dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")
-    docs = [
+    primary_docs = [
         Document(
             page_content="\n".join([
                 f"Name: {guest['name']}",
@@ -44,8 +80,18 @@ def load_guest_dataset():
                 f"Description: {guest['description']}",
                 f"Email: {guest['email']}"
             ]),
-            metadata={"name": guest["name"]}
+            metadata={"name": guest["name"], "source": "unit3-invitees"}
         )
         for guest in guest_dataset
     ]
-    return GuestInfoRetrieverTool(docs)
+    
+    # Secondary dataset (example: Wikipedia-like data)
+    secondary_docs = [
+        Document(
+            page_content="Name: Ada Lovelace\nDescription: Known as the first computer programmer, wrote the first algorithm for Charles Babbage's Analytical Engine.",
+            metadata={"name": "Ada Lovelace", "source": "wikipedia"}
+        )
+        # Add more secondary documents as needed
+    ]
+    
+    return MultiIndexRetrieverTool(primary_docs, secondary_docs)