File size: 4,183 Bytes
b81ac13 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
# Example: Copy tool implementations from sample_agent.tools here
# Math tools
def multiply(a: int, b: int) -> int:
"""Multiply two numbers.
Args:
a: first int
b: second int
"""
return a * b
def add(a: int, b: int) -> int:
"""Add two numbers.
Args:
a: first int
b: second int
"""
return a + b
def subtract(a: int, b: int) -> int:
"""Subtract two numbers.
Args:
a: first int
b: second int
"""
return a - b
def divide(a: int, b: int) -> float:
"""Divide two numbers.
Args:
a: first int
b: second int
"""
if b == 0:
raise ValueError("Cannot divide by zero.")
return a / b
def modulus(a: int, b: int) -> int:
"""Get the modulus of two numbers.
Args:
a: first int
b: second int
"""
return a % b
# Wikipedia search tool
def wiki_search(query: str) -> str:
"""Search Wikipedia for a query and return maximum 2 results.
Args:
query: The search query."""
try:
from langchain_community.document_loaders import WikipediaLoader
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 formatted_search_docs
except Exception as e:
return f"Error in wiki_search: {e}"
# Web search tool
def web_search(query: str) -> str:
"""Search Tavily for a query and return maximum 3 results.
Args:
query: The search query."""
try:
from langchain_community.tools.tavily_search import TavilySearchResults
search_tool = TavilySearchResults(max_results=3)
search_docs = search_tool.invoke({"query": query})
# Each doc is a dict, not an object with .metadata/.page_content
formatted_search_docs = "\n\n---\n\n".join(
[
f'<Document source="{doc.get("source", "")}" page="{doc.get("page", "")}"/>\n{doc.get("content", "")}\n</Document>'
for doc in search_docs
])
return formatted_search_docs
except Exception as e:
return f"Error in web_search: {e}"
# Arxiv search tool
def arvix_search(query: str) -> str:
"""Search Arxiv for a query and return maximum 3 result.
Args:
query: The search query."""
try:
from langchain_community.document_loaders import ArxivLoader
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 formatted_search_docs
except Exception as e:
return f"Error in arvix_search: {e}"
TOOL_REGISTRY = {
"multiply": {
"description": "Multiply two numbers. Usage: multiply(a, b)",
"function": multiply,
},
"add": {
"description": "Add two numbers. Usage: add(a, b)",
"function": add,
},
"subtract": {
"description": "Subtract two numbers. Usage: subtract(a, b)",
"function": subtract,
},
"divide": {
"description": "Divide two numbers. Usage: divide(a, b)",
"function": divide,
},
"modulus": {
"description": "Get the modulus of two numbers. Usage: modulus(a, b)",
"function": modulus,
},
"wiki_search": {
"description": "Search Wikipedia for a query and return up to 2 results. Usage: wiki_search(query)",
"function": wiki_search,
},
"web_search": {
"description": "Search Tavily for a query and return up to 3 results. Usage: web_search(query)",
"function": web_search,
},
"arvix_search": {
"description": "Search Arxiv for a query and return up to 3 results. Usage: arvix_search(query)",
"function": arvix_search,
},
}
|