RAG tool for guest story

.gitignore

@@ -1,2 +1,3 @@
-.venv
-.env
+.venv/
+.env
+__pycache__/

app.py

@@ -1,7 +1,18 @@
 import gradio as gr
+from smolagents import CodeAgent, HfApiModel
+from retriever import load_guest_dataset
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Initialize the Hugging Face model
+model = HfApiModel()
+guest_info_tool = load_guest_dataset()
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+# Create Alfred, our gala agent, with the guest info tool
+alfred = CodeAgent(tools=[guest_info_tool], model=model)
+
+
+def query_guest(guest_name):
+  return alfred.run(f"Tell me about our guest named '{guest_name}'.")
+
+
+demo = gr.Interface(fn=query_guest, inputs="text", outputs="text")
 demo.launch()

retriever.py

@@ -0,0 +1,48 @@
+import datasets
+from langchain.docstore.document import Document
+
+from smolagents import Tool
+from langchain_community.retrievers import BM25Retriever
+
+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.invoke(query)
+        if results:
+            return "\n\n".join([doc.page_content for doc in results[:3]])
+        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 GuestInfoRetrieverTool(docs)