Adding Tools

agent.py

@@ -1,11 +1,66 @@
 import os
 from dotenv import load_dotenv
 from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_core.tools import tool
+from supabase.client import Client, create_client
 
 load_dotenv()
 
+
+@tool
+def wikipedia_search(query: str) -> str:
+    """Search Wikipedia for a query and return maximum 2 results
+    Args:
+        query: The search string
+    """
+    docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    all_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 docs
+        ]
+    )
+    return {"wikipedia_results": all_search_docs}
+
+
+@tool
+def web_search(query: str) -> str:
+    """Search Tavily for a query and return maximum 3 results.
+
+    Args:
+        query: The search query."""
+    docs = TavilySearchResults(max_results=3).invoke(query=query)
+    all_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 docs
+        ]
+    )
+    return {"web_results": all_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}
+
+
 with open("system_prompt.txt", "r", encoding="utf-8") as f:
     system_prompt = f.read()
 
 sys_msg = SystemMessage(system_prompt)
-print(sys_msg)
+supabase: Client = create_client(
+    os.environ.get("SUPABASE_URL"), os.environ.get("SUPABASE_SERVICE_KEY")
+)