improving retriever BM25 - using sentence transformers - adding memory management tool - combine with websearch - integrate multiple indexes

app.py

@@ -1,13 +1,29 @@
 import gradio as gr
 import random
-from smolagents import GradioUI, CodeAgent, HfApiModel
+from typing import TypedDict, Annotated
 
-# Import our custom tools from their modules
-from tools import DuckDuckGoSearchTool, WeatherInfoTool, HubStatsTool
+from langgraph.graph.message import add_messages
+from langchain_core.messages import AnyMessage, HumanMessage, AIMessage
+from langgraph.prebuilt import ToolNode
+from langgraph.graph import START, StateGraph
+from langgraph.prebuilt import tools_condition
+from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
+
+from tools import DuckDuckGoSearchTool, WeatherInfoTool, HubStatsTool, MemoryManagementTool
 from retriever import load_guest_dataset
 
 # Initialize the Hugging Face model
-model = HfApiModel()
+llm = HuggingFaceEndpoint(
+    repo_id="Qwen/Qwen2.5-Coder-32B-Instruct",
+    huggingfacehub_api_token=HF_TOKEN,
+)
+
+chat = ChatHuggingFace(llm=llm, verbose=True)
+tools = [guest_info_tool, weather_info_tool, hub_stats_tool, search_tool, memory_management_tool]
+chat_with_tools = chat.bind_tools(tools)
+
+# Initialize the memory store
+memory_management_tool = MemoryManagementTool()
 
 # Initialize the web search tool
 search_tool = DuckDuckGoSearchTool()
@@ -21,13 +37,47 @@ hub_stats_tool = HubStatsTool()
 # Load the guest dataset and initialize the guest info tool
 guest_info_tool = load_guest_dataset()
 
-# 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
-    planning_interval=3   # Enable planning every 3 steps
+
+# Generate the AgentState and Agent graph
+class AgentState(TypedDict):
+    messages: Annotated[list[AnyMessage], add_messages]
+
+def assistant(state: AgentState):
+    # Retrieve past messages from memory
+    past_messages = memory_store.retrieve(state["messages"])
+    # Add new messages to memory
+    memory_store.add(state["messages"])
+    
+    # Check if the query is about an unfamiliar guest
+    if "guest" in state["messages"][-1].content:
+        search_results = search_tool.forward(state["messages"][-1].content)
+        return {
+            "messages": [search_results]
+        }
+    else:
+        return {
+            "messages": [chat_with_tools.invoke(state["messages"] + past_messages)],
+        }
+
+## The graph
+builder = StateGraph(AgentState)
+
+# Define nodes: these do the work
+builder.add_node("assistant", assistant)
+builder.add_node("tools", ToolNode(tools))
+
+# Define the graph
+
+builder.add_edge(START, "assistant")
+builder.add_conditional_edges(
+    "assistant",
+    # If the latest message requires a tool, route to tools
+    # Otherwise, provide a direct response
+    tools_condition,
 )
 
+builder.add_edge("tools", "assistant")
+alfred = builder.compile()
+
 if __name__ == "__main__":
     GradioUI(alfred).launch()

retriever.py

@@ -1,7 +1,9 @@
 from smolagents import Tool
-from langchain_community.retrievers import BM25Retriever
+# from langchain_community.retrievers import BM25Retriever
 from langchain.docstore.document import Document
 import datasets
+from sentence_transformers import SentenceTransformer
+import torch
 
 
 class GuestInfoRetrieverTool(Tool):
@@ -17,13 +19,20 @@ class GuestInfoRetrieverTool(Tool):
 
     def __init__(self, docs):
         self.is_initialized = False
-        self.retriever = BM25Retriever.from_documents(docs)
-       
+        # 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):
-        results = self.retriever.get_relevant_documents(query)
+        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[:3]])
+            return "\n\n".join([doc.page_content for doc in results])
         else:
             return "No matching guest information found."
 

tools.py

@@ -2,11 +2,13 @@ from smolagents import DuckDuckGoSearchTool
 from smolagents import Tool
 import random
 from huggingface_hub import list_models
-
+from langgraph.store.memory import InMemoryStore
 
 # Initialize the DuckDuckGo search tool
-#search_tool = DuckDuckGoSearchTool()
-
+search_tool = DuckDuckGoSearchTool()
+store = InMemoryStore(
+    index={"embed": "openai:text-embedding-3-small"}
+)
 
 class WeatherInfoTool(Tool):
     name = "weather_info"
@@ -54,3 +56,23 @@ class HubStatsTool(Tool):
         except Exception as e:
             return f"Error fetching models for {author}: {str(e)}"
 
+class MemoryManagementTool(Tool):
+    name = "memory_management"
+    description = "Manages and queries conversation memory."
+    inputs = {
+        "query": {
+            "type": "string",
+            "description": "The query to search in memory."
+        }
+    }
+    output_type = "string"
+
+    def forward(self, query: str):
+        # Retrieve relevant memory entries
+        results = store.retrieve(query)
+        if results:
+            return "\n\n".join(results)
+        else:
+            return "No relevant memory found."
+
+