Update retriever.py

retriever.py

@@ -1,59 +1,56 @@
 from smolagents import Tool
 import datasets
-from langchain_core.documents import Document
 
-print("🔄 Loading RAG dataset...")
+print("📂 Loading guest dataset...")
 
-# Load the course dataset
 try:
+    # Load dataset as shown in course
     guest_dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")
     print(f"✅ Loaded {len(guest_dataset)} guests")
 except:
     # Fallback data
     guest_dataset = [
-        {"name": "Lady Ada Lovelace", "relation": "mathematician", "description": "First computer programmer", "email": "ada@example.com"},
-        {"name": "Dr. Nikola Tesla", "relation": "inventor", "description": "Electrical engineering pioneer", "email": "tesla@example.com"},
-        {"name": "Marie Curie", "relation": "scientist", "description": "Nobel prize winning physicist", "email": "marie@example.com"}
+        {"name": "Lady Ada Lovelace", "relation": "mathematician", 
+         "description": "First computer programmer", "email": "ada@example.com"},
+        {"name": "Dr. Nikola Tesla", "relation": "inventor", 
+         "description": "Electrical engineering pioneer", "email": "tesla@example.com"}
     ]
     print(f"⚠️ Using {len(guest_dataset)} sample guests")
 
-# Create documents for RAG
-docs = [
-    Document(
-        page_content=f"Name: {guest['name']}\nRelation: {guest['relation']}\nDescription: {guest['description']}\nEmail: {guest['email']}",
-        metadata={"name": guest["name"]}
-    )
-    for guest in guest_dataset
-]
-
 class GuestInfoRetrieverTool(Tool):
     name = "guest_info_retriever"
-    description = "Retrieves detailed information about gala guests using RAG."
+    description = "Retrieves detailed information about gala guests based on their name or relation."
     inputs = {
         "query": {
             "type": "string",
-            "description": "The name or relation of the guest."
+            "description": "The name or relation of the guest you want information about."
         }
     }
     output_type = "string"
 
     def forward(self, query: str):
-        # Simple RAG: search through documents
         query_lower = query.lower()
         results = []
         
-        for doc in docs:
-            if query_lower in doc.page_content.lower():
-                results.append(doc.page_content)
+        for guest in guest_dataset:
+            if (query_lower in guest['name'].lower() or 
+                query_lower in guest['relation'].lower()):
+                
+                guest_info = f"""Name: {guest['name']}
+Relation: {guest['relation']}
+Description: {guest['description']}
+Email: {guest['email']}"""
+                
+                results.append(guest_info)
                 if len(results) >= 3:
                     break
         
         if results:
-            return "🧠 RAG Results:\n\n" + "\n\n---\n\n".join(results)
+            return "\n\n---\n\n".join(results)
         else:
-            return f"No guest found for '{query}'"
+            return f"No guest found matching '{query}'"
 
-# Create RAG tool
+# Create tool instance
 guest_info_tool = GuestInfoRetrieverTool()
 print("✅ RAG tool created")