Update agent.py

agent.py

@@ -48,7 +48,7 @@ def subtract(a: int, b: int) -> int:
     return a - b
 
 @tool
-def divide(a: int, b: int) -> float:
+def divide(a: int, b: int) -> int:
     """Divide two numbers.
     
     Args:
@@ -70,7 +70,7 @@ def modulus(a: int, b: int) -> int:
     return a % b
 
 @tool
-def wiki_search(query: str) -> dict:
+def wiki_search(query: str) -> str:
     """Search Wikipedia for a query and return maximum 2 results.
     
     Args:
@@ -84,7 +84,7 @@ def wiki_search(query: str) -> dict:
     return {"wiki_results": formatted_search_docs}
 
 @tool
-def web_search(query: str) -> dict:
+def web_search(query: str) -> str:
     """Search Tavily for a query and return maximum 3 results.
     
     Args:
@@ -98,7 +98,7 @@ def web_search(query: str) -> dict:
     return {"web_results": formatted_search_docs}
 
 @tool
-def arxiv_search(query: str) -> dict:
+def arvix_search(query: str) -> str:
     """Search Arxiv for a query and return maximum 3 result.
     
     Args:
@@ -109,7 +109,7 @@ def arxiv_search(query: str) -> dict:
             f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
             for doc in search_docs
         ])
-    return {"arxiv_results": formatted_search_docs}
+    return {"arvix_results": formatted_search_docs}
 
 
 
@@ -147,8 +147,7 @@ tools = [
     modulus,
     wiki_search,
     web_search,
-    arxiv_search,
-    create_retriever_tool,
+    arvix_search,
 ]
 
 # Build graph function
@@ -206,10 +205,9 @@ 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="google")
+    graph = build_graph(provider="groq")
     # Run the graph
     messages = [HumanMessage(content=question)]
     messages = graph.invoke({"messages": messages})
     for m in messages["messages"]:
         m.pretty_print()
-