Update retriever.py

retriever.py

@@ -4,9 +4,7 @@ from langchain.docstore.document import Document
 import datasets
 import pandas as pd
 import os
-import logging
 
-logging.basicConfig(level=logging.INFO)
 
 class GuestInfoRetrieverTool(Tool):
     name = "guest_info_retriever"
@@ -19,10 +17,16 @@ class GuestInfoRetrieverTool(Tool):
     }
     output_type = "string"
 
+    #def __init__(self, docs):
+        #self.is_initialized = False
+        #self.retriever = BM25Retriever.from_documents(docs)
+
     def __init__(self, docs):
-        self.docs = docs
+        self.is_initialized = False
+        self.docs = docs  # 🔁 store the original list manually
         self.retriever = BM25Retriever.from_documents(docs)
 
+
     def _generate_conversation_starter(self, doc: Document):
         lines = doc.page_content.splitlines()
         name = None
@@ -43,22 +47,24 @@ class GuestInfoRetrieverTool(Tool):
         else:
             return "Try asking about their background—it sounds fascinating!"
 
-    def forward(self, query: str):
-        query = query.strip().lower()
+    #def forward(self, query: str):
+        # Handle special case for full guest listing
+        #if "list" in query.lower() and "guest" in query.lower() and "name" in query.lower():
+            #return "\n".join([
+                #doc.metadata.get("name", "Unknown") for doc in self.retriever.docs
+                #])
 
-        # Handle guest listing queries
-        guest_list_keywords = [
-            "list guests", "guest names", "list all guests",
-            "show guests", "all guests", "everyone invited"
-        ]
-        if any(keyword in query for keyword in guest_list_keywords):
-            logging.info("Listing all guests from provided dataset.")
+    def forward(self, query: str):
+        if any(keyword in query.lower() for keyword in ["list guests", "guest names", "list all guests", "show guests", "all guests", "everyone invited"]):
+        #if "list" in query.lower() and "guest" in query.lower() and "name" in query.lower():
             return "\n".join([doc.metadata.get("name", "Unknown") for doc in self.docs])
 
-        # Fallback to BM25 search
+        
+    # Default BM25 retrieval
         results = self.retriever.get_relevant_documents(query)
         if results:
             responses = []
+            #for doc in results[:3]:
             for doc in results[:10]:
                 content = doc.page_content
                 starter = self._generate_conversation_starter(doc)
@@ -69,14 +75,13 @@ class GuestInfoRetrieverTool(Tool):
 
 
 def load_guest_dataset(file_path: str = None, show_example: bool = True):
+    
     """
-    Loads a guest dataset from a CSV/JSON file or from Hugging Face if no file provided.
-    Ensures necessary columns exist.
-    Returns a Tool that can search for guest info.
+    Loads guest dataset either from a file (CSV/JSON) or the Hugging Face default dataset.
+    If using the Hugging Face dataset, optionally prints a preview example.
     """
     if file_path and os.path.exists(file_path):
         ext = os.path.splitext(file_path)[1].lower()
-        logging.info(f"📁 Loading guest data from: {file_path}")
         if ext == ".csv":
             df = pd.read_csv(file_path)
         elif ext == ".json":
@@ -84,17 +89,12 @@ def load_guest_dataset(file_path: str = None, show_example: bool = True):
         else:
             raise ValueError("Unsupported file format. Use .csv or .json.")
     else:
-        logging.info("📡 Loading default guest data from Hugging Face.")
         guest_dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")
         df = pd.DataFrame(guest_dataset)
         if show_example:
             print("\n📌 Example guest from Hugging Face dataset:\n")
             print(df.head(1).to_markdown(index=False))
 
-    required_columns = {"name", "relation", "description", "email"}
-    if not required_columns.issubset(df.columns):
-        raise ValueError(f"Missing required columns. Expected: {required_columns}")
-
     docs = [
         Document(
             page_content="\n".join([
@@ -108,5 +108,4 @@ def load_guest_dataset(file_path: str = None, show_example: bool = True):
         for _, row in df.iterrows()
     ]
 
-    logging.info(f"✅ Loaded {len(docs)} guests into retrieval tool.")
-    return GuestInfoRetrieverTool(docs)
+    return GuestInfoRetrieverTool(docs)