Update agent.py

agent.py

@@ -5,111 +5,99 @@ from langgraph.prebuilt import ToolNode, tools_condition
 
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_groq import ChatGroq
-from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
 
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
+from langchain_community.vectorstores import SupabaseVectorStore
 
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
 
-# ---- Basic Arithmetic Utilities as Classes ---- #
-
-class MultiplyTool:
-    name = "multiply"
-    description = "Returns the product of two integers."
-
-    def __call__(self, a: int, b: int) -> int:
-        return a * b
-
-class AddTool:
-    name = "add"
-    description = "Returns the sum of two integers."
-
-    def __call__(self, a: int, b: int) -> int:
-        return a + b
-
-class SubtractTool:
-    name = "subtract"
-    description = "Returns the difference between two integers."
-
-    def __call__(self, a: int, b: int) -> int:
-        return a - b
-
-class DivideTool:
-    name = "divide"
-    description = "Performs division and handles zero division errors."
-
-    def __call__(self, a: int, b: int) -> float:
-        if b == 0:
-            raise ValueError("Division by zero is undefined.")
-        return a / b
-
-class ModulusTool:
-    name = "modulus"
-    description = "Returns the remainder after division."
-
-    def __call__(self, a: int, b: int) -> int:
-        return a % b
-
-# ---- Search Tools as Classes ---- #
-
-class SearchWikipediaTool:
-    name = "search_wikipedia"
-    description = "Search Wikipedia for a query and return maximum 2 results."
-
-    def __call__(self, query: str) -> str:
-        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}
-
-class SearchWebTool:
-    name = "search_web"
-    description = "Search Tavily for a query and return maximum 3 results."
-
-    def __call__(self, query: str) -> str:
-        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}
-
-class SearchArxivTool:
-    name = "search_arxiv"
-    description = "Search Arxiv for a query and return maximum 3 results."
-
-    def __call__(self, query: str) -> str:
-        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}
+from supabase.client import create_client, Client
+
+
+# Load environment variables
+
+# ---- Basic Arithmetic Utilities ---- #
+@tool
+def multiply(a: int, b: int) -> int:
+    """Returns the product of two integers."""
+    return a * b
+
+@tool
+def add(a: int, b: int) -> int:
+    """Returns the sum of two integers."""
+    return a + b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """Returns the difference between two integers."""
+    return a - b
+
+@tool
+def divide(a: int, b: int) -> float:
+    """Performs division and handles zero division errors."""
+    if b == 0:
+        raise ValueError("Division by zero is undefined.")
+    return a / b
+
+@tool
+def modulus(a: int, b: int) -> int:
+    """Returns the remainder after division."""
+    return a % b
+
+
+# ---- Search Tools ---- #
+@tool
+def search_wikipedia(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 search_web(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 search_arxiv(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}
+
 
 system_message = SystemMessage(content="""You are a helpful assistant tasked with answering questions using a set of tools. 
 Now, I will ask you a question. Report your thoughts, and finish your answer with the following template: 
 FINAL ANSWER: [YOUR FINAL ANSWER]. 
 YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
-Your answer should only start with \"FINAL ANSWER: \" then follows with the answer. """)
-
-# Instantiate tools
-multiply = MultiplyTool()
-add = AddTool()
-subtract = SubtractTool()
-divide = DivideTool()
-modulus = ModulusTool()
-search_wikipedia = SearchWikipediaTool()
-search_web = SearchWebTool()
-search_arxiv = SearchArxivTool()
-
-# Toolset
+Your answer should only start with "FINAL ANSWER: ", then follows with the answer. """)
+
 toolset = [
     multiply,
     add,
@@ -170,3 +158,4 @@ if __name__ == "__main__":
     # Print the final output
     for m in output_state["messages"]:
         print(m.content)
+