Implemented RAG

app.py

@@ -1,10 +1,10 @@
-import gradio as gr
-import random
+import chromadb
+from chromadb.utils import embedding_functions
 from smolagents import GradioUI, CodeAgent, HfApiModel
 
+from retriever import load_guest_dataset
 # Import our custom tools from their modules
 from tools import DuckDuckGoSearchTool, WeatherInfoTool, HubStatsTool
-from retriever import load_guest_dataset
 
 # Initialize the Hugging Face model
 model = HfApiModel()
@@ -19,15 +19,19 @@ weather_info_tool = WeatherInfoTool()
 hub_stats_tool = HubStatsTool()
 
 # Load the guest dataset and initialize the guest info tool
-guest_info_tool = load_guest_dataset()
+chroma_client = chromadb.Client()
+default_ef = embedding_functions.DefaultEmbeddingFunction()
+collection = chroma_client.create_collection(name="guests-collection", embedding_function=default_ef)
+
+guest_info_tool = load_guest_dataset(collection)
 
 # Create Alfred with all the tools
 alfred = CodeAgent(
     tools=[guest_info_tool, weather_info_tool, hub_stats_tool, search_tool], 
     model=model,
-    add_base_tools=True,  # Add any additional base tools
+    add_base_tools=False,  # Add any additional base tools
     planning_interval=3   # Enable planning every 3 steps
 )
 
 if __name__ == "__main__":
-    GradioUI(alfred).launch()
+    GradioUI(alfred).launch()

requirements.txt

@@ -1,4 +1,4 @@
 datasets 
 smolagents 
-langchain-community 
-rank_bm25
+langchain-community
+langchain-huggingface

retriever.py

@@ -1,7 +1,7 @@
-from smolagents import Tool
-from langchain_community.retrievers import BM25Retriever
-from langchain.docstore.document import Document
+from uuid import uuid4
+
 import datasets
+from smolagents import Tool
 
 
 class GuestInfoRetrieverTool(Tool):
@@ -15,39 +15,38 @@ class GuestInfoRetrieverTool(Tool):
     }
     output_type = "string"
 
-    def __init__(self, docs):
+    def __init__(self, vector_store):
         self.is_initialized = False
-        self.retriever = BM25Retriever.from_documents(docs)
-       
+        self.vector_store = vector_store
 
     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."
+        result = self.vector_store.query(
+            query_texts=[query],
+            n_results=3
+        )
+
+        distances = [distance for distance in result['distances'][0] if distance < 1.3]
+        docs = result['documents'][0]
+        return "\n\n".join([docs[idx] for idx in range(0, len(distances))])
 
 
-def load_guest_dataset():
+def load_guest_dataset(vector_store):
     # 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
-    ]
+    for guest in guest_dataset:
+        vector_store.add(
+            documents=[
+                "\n".join([
+                    f"Name: {guest['name']}",
+                    f"Relation: {guest['relation']}",
+                    f"Description: {guest['description']}",
+                    f"Email: {guest['email']}"
+                ])
+            ],
+            metadatas=[{"name": guest["name"]}],
+            ids=[str(uuid4())])
 
     # Return the tool
-    return GuestInfoRetrieverTool(docs)
-
-
-
+    return GuestInfoRetrieverTool(vector_store)