Update tools.py
Browse files
tools.py
CHANGED
|
@@ -0,0 +1,52 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tool definitions for LangGraph agent."""
|
| 2 |
+
from langchain_core.tools import tool
|
| 3 |
+
from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
|
| 4 |
+
from langchain_community.tools.tavily_search import TavilySearchResults
|
| 5 |
+
|
| 6 |
+
@tool
|
| 7 |
+
def multiply(a: int, b: int) -> int:
|
| 8 |
+
"""Multiply two numbers."""
|
| 9 |
+
return a * b
|
| 10 |
+
|
| 11 |
+
@tool
|
| 12 |
+
def add(a: int, b: int) -> int:
|
| 13 |
+
"""Add two numbers."""
|
| 14 |
+
return a + b
|
| 15 |
+
|
| 16 |
+
@tool
|
| 17 |
+
def subtract(a: int, b: int) -> int:
|
| 18 |
+
"""Subtract two numbers."""
|
| 19 |
+
return a - b
|
| 20 |
+
|
| 21 |
+
@tool
|
| 22 |
+
def divide(a: int, b: int) -> float:
|
| 23 |
+
"""Divide two numbers."""
|
| 24 |
+
if b == 0:
|
| 25 |
+
raise ValueError("Cannot divide by zero.")
|
| 26 |
+
return a / b
|
| 27 |
+
|
| 28 |
+
@tool
|
| 29 |
+
def modulus(a: int, b: int) -> int:
|
| 30 |
+
"""Get modulus of two numbers."""
|
| 31 |
+
return a % b
|
| 32 |
+
|
| 33 |
+
@tool
|
| 34 |
+
def wiki_search(query: str) -> str:
|
| 35 |
+
"""Search Wikipedia for a query and return top results."""
|
| 36 |
+
docs = WikipediaLoader(query=query, load_max_docs=2).load()
|
| 37 |
+
return "\n\n".join(doc.page_content for doc in docs)
|
| 38 |
+
|
| 39 |
+
@tool
|
| 40 |
+
def arxiv_search(query: str) -> str:
|
| 41 |
+
"""Search Arxiv for a query and return top results."""
|
| 42 |
+
docs = ArxivLoader(query=query, load_max_docs=2).load()
|
| 43 |
+
return "\n\n".join(doc.page_content[:1000] for doc in docs)
|
| 44 |
+
|
| 45 |
+
@tool
|
| 46 |
+
def web_search(query: str) -> str:
|
| 47 |
+
"""Search the web using Tavily."""
|
| 48 |
+
docs = TavilySearchResults(max_results=3).invoke(query=query)
|
| 49 |
+
return "\n\n".join(doc.page_content for doc in docs)
|
| 50 |
+
|
| 51 |
+
# Export list of all tools
|
| 52 |
+
TOOLS = [multiply, add, subtract, divide, modulus, wiki_search, arxiv_search, web_search]
|