Testing a simpler agent

agents/agent.py

@@ -1,80 +1,162 @@
-"""LangGraph Agent for GAIA Benchmark"""
-from tools.SearchToolkit import wiki_search, web_search, arxiv_search, vector_store
-from tools.MathsToolkit import (
-    multiply, add, subtract, divide, modulus, power, square_root
-)
-from tools.ImagesToolkit import (
-    analyze_image,
-    transform_image,
-    draw_on_image,
-    generate_simple_image,
-    combine_images
-)
-from tools.DocumentsToolkit import (
-    save_and_read_file,
-    download_file_from_url,
-    extract_text_from_image,
-    analyze_csv_file,
-    analyze_excel_file,
-    analyze_word_file,
-    analyze_pdf_file
-)
-from tools.CodeToolkit import execute_code_multilang
-from langchain_groq import ChatGroq
-from langchain_core.messages import SystemMessage, HumanMessage
-from langgraph.prebuilt import tools_condition, ToolNode
-from langgraph.graph import START, StateGraph, MessagesState
+"""LangGraph Agent"""
 import os
 from dotenv import load_dotenv
-# Load environment variables
+from langgraph.graph import START, StateGraph, MessagesState
+from langgraph.prebuilt import tools_condition
+from langgraph.prebuilt import ToolNode
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_groq import ChatGroq
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
+from langchain_community.document_loaders import ArxivLoader
+from langchain_community.vectorstores import SupabaseVectorStore
+from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_core.tools import tool
+from langchain.tools.retriever import create_retriever_tool
+from supabase.client import Client, create_client
+
 load_dotenv()
 
-prompt_path = os.path.join(os.path.dirname(__file__), "../prompts")
 
+@tool
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers.
+    Args:
+        a: first int
+        b: second int
+    """
+    return a * b
+
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add two numbers.
+
+    Args:
+        a: first int
+        b: second int
+    """
+    return a + b
+
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract two numbers.
+
+    Args:
+        a: first int
+        b: second int
+    """
+    return a - b
+
+
+@tool
+def divide(a: int, b: int) -> int:
+    """Divide two numbers.
+
+    Args:
+        a: first int
+        b: second int
+    """
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
+
+@tool
+def modulus(a: int, b: int) -> int:
+    """Get the modulus of two numbers.
+
+    Args:
+        a: first int
+        b: second int
+    """
+    return a % b
 
-# Load system prompt
-with open(os.path.join(prompt_path, "system_prompt.txt"), "r", encoding="utf-8") as f:
+
+@tool
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return maximum 2 results.
+
+    Args:
+        query: The search query."""
+    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"wiki_results": formatted_search_docs}
+
+
+@tool
+def web_search(query: str) -> str:
+    """Search Tavily for a query and return maximum 3 results.
+
+    Args:
+        query: The search query."""
+    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"web_results": formatted_search_docs}
+
+
+@tool
+def arvix_search(query: str) -> str:
+    """Search Arxiv for a query and return maximum 3 result.
+
+    Args:
+        query: The search query."""
+    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"arvix_results": formatted_search_docs}
+
+
+# load the system prompt from the file
+with open("system_prompt.txt", "r", encoding="utf-8") as f:
     system_prompt = f.read()
+
+# System message
 sys_msg = SystemMessage(content=system_prompt)
 
-# Toolset
-tools = [
-    # SearchToolkit
-    web_search,
-    wiki_search,
-    arxiv_search,
+# build a retriever
+embeddings = HuggingFaceEmbeddings(
+    model_name="sentence-transformers/all-mpnet-base-v2")  # dim=768
+supabase: Client = create_client(
+    os.environ.get("SUPABASE_URL"),
+    os.environ.get("SUPABASE_SERVICE_KEY"))
+vector_store = SupabaseVectorStore(
+    client=supabase,
+    embedding=embeddings,
+    table_name="documents",
+    query_name="match_documents_langchain",
+)
+create_retriever_tool = create_retriever_tool(
+    retriever=vector_store.as_retriever(),
+    name="Question Search",
+    description="A tool to retrieve similar questions from a vector store.",
+)
 
-    # MathsToolkit
+
+tools = [
     multiply,
     add,
     subtract,
     divide,
     modulus,
-    power,
-    square_root,
-
-    # DocumentsToolkit
-    save_and_read_file,
-    download_file_from_url,
-    extract_text_from_image,
-    analyze_csv_file,
-    analyze_excel_file,
-    analyze_word_file,
-    analyze_pdf_file,
-
-    # CodeToolkit
-    execute_code_multilang,
-
-    # ImagesToolkit
-    analyze_image,
-    transform_image,
-    draw_on_image,
-    generate_simple_image,
-    combine_images,
+    wiki_search,
+    web_search,
+    arvix_search,
 ]
 
-# Build LangGraph workflow
-
 
 def build_graph():
     llm = ChatGroq(model="qwen-qwq-32b", temperature=0)
@@ -84,14 +166,13 @@ def build_graph():
         return {"messages": [llm_with_tools.invoke(state["messages"])]}
 
     def retriever(state: MessagesState):
-        similar = vector_store.similarity_search(state["messages"][0].content)
-        if similar:
-            reference = HumanMessage(
-                content=f"Here is a similar Q&A that might help: \n\n{similar[0].page_content}"
-            )
-            return {"messages": [sys_msg] + state["messages"] + [reference]}
-        else:
-            return {"messages": [sys_msg] + state["messages"]}
+        """Retriever node"""
+        similar_question = vector_store.similarity_search(
+            state["messages"][0].content)
+        example_msg = HumanMessage(
+            content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
+        )
+        return {"messages": [sys_msg] + state["messages"] + [example_msg]}
 
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever)