Update agent.py

agent.py

@@ -1,4 +1,3 @@
-
 """LangGraph Agent"""
 import os
 from dotenv import load_dotenv
@@ -49,7 +48,7 @@ def subtract(a: int, b: int) -> int:
     return a - b
 
 @tool
-def divide(a: int, b: int) -> int:
+def divide(a: int, b: int) -> float:
     """Divide two numbers.
     
     Args:
@@ -71,7 +70,7 @@ def modulus(a: int, b: int) -> int:
     return a % b
 
 @tool
-def wiki_search(query: str) -> str:
+def wiki_search(query: str) -> dict:
     """Search Wikipedia for a query and return maximum 2 results.
     
     Args:
@@ -85,7 +84,7 @@ def wiki_search(query: str) -> str:
     return {"wiki_results": formatted_search_docs}
 
 @tool
-def web_search(query: str) -> str:
+def web_search(query: str) -> dict:
     """Search Tavily for a query and return maximum 3 results.
     
     Args:
@@ -99,7 +98,7 @@ def web_search(query: str) -> str:
     return {"web_results": formatted_search_docs}
 
 @tool
-def arvix_search(query: str) -> str:
+def arxiv_search(query: str) -> dict:
     """Search Arxiv for a query and return maximum 3 result.
     
     Args:
@@ -110,7 +109,7 @@ def arvix_search(query: str) -> str:
             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}
+    return {"arxiv_results": formatted_search_docs}
 
 
 
@@ -148,7 +147,8 @@ tools = [
     modulus,
     wiki_search,
     web_search,
-    arvix_search,
+    arxiv_search,
+    create_retriever_tool,
 ]
 
 # Build graph function
@@ -206,7 +206,7 @@ def build_graph(provider: str = "groq"):
 if __name__ == "__main__":
     question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
     # Build the graph
-    graph = build_graph(provider="groq")
+    graph = build_graph(provider="google")
     # Run the graph
     messages = [HumanMessage(content=question)]
     messages = graph.invoke({"messages": messages})