feat: add handoff tool for agent creation and embedding query tool for database

tools/__init__.py

@@ -1,5 +1,8 @@
+from tools.handoff_tools import create_handoff_tool
 from tools.math_tools import add, div, mod, mult, sub
-from tools.search_tools import arxiv_search, internet_search, wiki_search
+from tools.search_tools import arxiv_search
+from tools.search_tools import create_retriever_from_supabase as retriever_tool
+from tools.search_tools import internet_search, wiki_search
 
 __all__ = [
     "add",
@@ -8,6 +11,8 @@ __all__ = [
     "mult",
     "sub",
     "arxiv_search",
+    "retriever_tool",
     "internet_search",
     "wiki_search",
+    "create_handoff_tool",
 ]

tools/handoff_tools.py

@@ -0,0 +1,31 @@
+from typing import Annotated
+
+from langchain_core.tools import InjectedToolCallId, tool
+from langgraph.graph import MessagesState
+from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+
+
+def create_handoff_tool(*, agent_name: str, description: str | None = None):
+    name = f"transfer_to_{agent_name}"
+    description = description or f"Ask {agent_name} for help."
+
+    @tool(name_or_callable=name, description=description)
+    def handoff_tool(
+        state: Annotated[MessagesState, InjectedState],
+        tool_call_id: Annotated[str, InjectedToolCallId],
+    ) -> Command:
+        tool_message = {
+            "role": "tool",
+            "content": f"Successfully transferred to {agent_name}",
+            "name": name,
+            "tool_call_id": tool_call_id,
+        }
+
+        return Command(
+            goto=agent_name,
+            update={**state, "messages": state["messages"] + [tool_message]},
+            graph=Command.PARENT,
+        )
+
+    return handoff_tool

tools/search_tools.py

@@ -1,8 +1,43 @@
+import os
 from typing import Dict, List
 
+from dotenv import load_dotenv
 from langchain_community.document_loaders import ArxivLoader, WikipediaLoader
+from langchain_community.vectorstores import SupabaseVectorStore
 from langchain_core.tools import tool
+from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_tavily import TavilySearch
+from supabase.client import create_client
+
+load_dotenv()
+
+
+@tool
+def create_retriever_from_supabase(query: str) -> str:
+    """Search for similar documents in the Supabase vector store.
+
+    This tool uses semantic search to find documents that are semantically similar to the provided query.
+    It leverages the Supabase vector store and HuggingFace embeddings to perform the search.
+
+    Args:
+        query (str): The search query to find similar documents.
+
+    Returns:
+        str: A list of documents that are semantically similar to the query.
+    """
+    embeddings = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-mpnet-base-v2"
+    )
+    supabase = create_client(os.environ["SUPABASE_URL"], os.environ["SUPABASE_KEY"])
+
+    vector_store = SupabaseVectorStore(
+        client=supabase,
+        embedding=embeddings,
+        table_name="documents",
+        query_name="match_documents_langchain",
+    )
+
+    return vector_store.similarity_search(query)
 
 
 @tool
@@ -10,7 +45,7 @@ def internet_search(query: str) -> Dict[str, List[Dict[str, str]]]:
     """Perform a web search using Tavily Search API.
 
     This tool searches the web for relevant information based on the provided query.
-    It returns up to 2 most relevant results with their sources, titles, and content.
+    It returns up to 3 most relevant results with their sources, titles, and content.
 
     Args:
         query (str): The search query to look up on the web.
@@ -22,7 +57,7 @@ def internet_search(query: str) -> Dict[str, List[Dict[str, str]]]:
             - Title: Title of the webpage
             - Content: Main content/text from the webpage
     """
-    response = TavilySearch(max_results=1).invoke(query)
+    response = TavilySearch(max_results=3).invoke(query)
 
     formatted_answer = [
         {
@@ -41,7 +76,7 @@ def wiki_search(query: str) -> Dict[str, List[Dict[str, str]]]:
     """Search Wikipedia articles using the provided query.
 
     This tool searches Wikipedia for articles matching the query and returns
-    up to 2 most relevant results with their sources, titles, and content.
+    up to 3 most relevant results with their sources, titles, and content.
 
     Args:
         query (str): The search query to look up on Wikipedia.
@@ -53,7 +88,7 @@ def wiki_search(query: str) -> Dict[str, List[Dict[str, str]]]:
             - Title: Title of the Wikipedia article
             - Content: Main content/text from the article
     """
-    docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    docs = WikipediaLoader(query=query, load_max_docs=3).load()
 
     formatted_answer = [
         {
@@ -72,7 +107,7 @@ def arxiv_search(query: str) -> Dict[str, List[Dict[str, str]]]:
     """Search academic papers on arXiv using the provided query.
 
     This tool searches arXiv for academic papers matching the query and returns
-    up to 2 most relevant results with their sources, titles, and content.
+    up to 3 most relevant results with their sources, titles, and content.
 
     Args:
         query (str): The search query to look up on arXiv.
@@ -84,7 +119,7 @@ def arxiv_search(query: str) -> Dict[str, List[Dict[str, str]]]:
             - Title: Title of the academic paper
             - Content: Main content/abstract of the paper
     """
-    docs = ArxivLoader(query=query, load_max_docs=2).load()
+    docs = ArxivLoader(query=query, load_max_docs=3).load()
 
     formatted_answer = [
         {