updating retriever to use langgraph

retriever.py

@@ -1,44 +1,11 @@
-from smolagents import Tool
-# from langchain_community.retrievers import BM25Retriever
+from langchain.tools import Tool
 from langchain.docstore.document import Document
-import datasets
 from sentence_transformers import SentenceTransformer
 import torch
+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
-        # Use sentence-transformers for embeddings
-        self.model = SentenceTransformer('all-MiniLM-L6-v2')
-        self.embeddings = self.model.encode([doc.page_content for doc in docs], convert_to_tensor=True)
-        self.docs = docs
-
-    def forward(self, query: str):
-        query_embedding = self.model.encode(query, convert_to_tensor=True)
-        # Compute cosine similarities
-        similarities = torch.nn.functional.cosine_similarity(query_embedding, self.embeddings)
-        # Get the top 3 most similar documents
-        top_k = torch.topk(similarities, k=3)
-        results = [self.docs[i] for i in top_k.indices]
-        if results:
-            return "\n\n".join([doc.page_content for doc in results])
-        else:
-            return "No matching guest information found."
-
-
+# Load the dataset
 def load_guest_dataset():
-    # Load the dataset
     guest_dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")
 
     # Convert dataset entries into Document objects
@@ -55,8 +22,30 @@ def load_guest_dataset():
         for guest in guest_dataset
     ]
 
-    # Return the tool
-    return GuestInfoRetrieverTool(docs)
+    # Initialize the sentence-transformers model
+    model = SentenceTransformer('all-MiniLM-L6-v2')
+    embeddings = model.encode([doc.page_content for doc in docs], convert_to_tensor=True)
+
+    # Define the extraction function
+    def extract_text(query: str) -> str:
+        """Retrieves detailed information about gala guests based on their name or relation."""
+        query_embedding = model.encode(query, convert_to_tensor=True)
+        similarities = torch.nn.functional.cosine_similarity(query_embedding, embeddings)
+        top_k = torch.topk(similarities, k=3)
+        results = [docs[i] for i in top_k.indices]
+        if results:
+            return "\n\n".join([doc.page_content for doc in results])
+        else:
+            return "No matching guest information found."
+
+    # Create the tool
+    guest_info_tool = Tool(
+        name="guest_info_retriever",
+        func=extract_text,
+        description="Retrieves detailed information about gala guests based on their name or relation."
+    )
+
+    return guest_info_tool