Create retriever.py

retriever.py

@@ -0,0 +1,84 @@
+from smolagents import Tool
+from langchain_community.retrievers import BM25Retriever
+from langchain.docstore.document import Document
+import datasets
+
+
+class GuestInfoRetrieverTool(Tool):
+    name = "guest_info_retriever"
+    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 you want information about."
+        }
+    }
+    output_type = "string"
+
+    def __init__(self, docs):
+        self.is_initialized = False
+        self.retriever = BM25Retriever.from_documents(docs)
+       
+
+    #def forward(self, query: str):
+      #   results = self.retriever.get_relevant_documents(query)
+       #  if results:
+        #     return "\n\n".join([doc.page_content for doc in results[:3]])
+       #  else:
+          #   return "No matching guest information found."
+
+
+    def _generate_conversation_starter(self, doc: Document):
+        lines = doc.page_content.splitlines()
+        name = None
+        description = ""
+        for line in lines:
+            if line.startswith("Name:"):
+                name = line.replace("Name:", "").strip()
+            if line.startswith("Description:"):
+                description = line.replace("Description:", "").strip()
+
+        # Example heuristic: use keywords from description
+        # You could expand this with keyword extraction or simple NLP parsing
+        interests = []
+        for interest in ["art", "science", "sports", "music", "history", "technology", "travel", "literature"]:
+            if interest.lower() in description.lower():
+                interests.append(interest)
+
+        if interests:
+            return f"A good icebreaker could be: 'I heard you're into {interests[0]}. What's your favorite part about it?'"
+        else:
+            return "Try asking about their background—it sounds fascinating!"
+
+    def forward(self, query: str):
+        results = self.retriever.get_relevant_documents(query)
+        if results:
+            responses = []
+            for doc in results[:3]:
+                content = doc.page_content
+                starter = self._generate_conversation_starter(doc)
+                responses.append(f"{content}\n\n{starter}")
+            return "\n\n---\n\n".join(responses)
+        else:
+            return "No matching guest information found."
+
+def load_guest_dataset():
+    # Load the dataset
+    guest_dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")
+
+    # Convert dataset entries into Document objects
+    docs = [
+        Document(
+            page_content="\n".join([
+                f"Name: {guest['name']}",
+                f"Relation: {guest['relation']}",
+                f"Description: {guest['description']}",
+                f"Email: {guest['email']}"
+            ]),
+            metadata={"name": guest["name"]}
+        )
+        for guest in guest_dataset
+    ]
+
+    # Return the tool
+    return GuestInfoRetrieverTool(docs)