File size: 4,986 Bytes
1e22d0b |
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 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 |
import cmath
import json
from langchain_community.document_loaders import ArxivLoader
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.tools import tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader
@tool
def wiki_search(query: str) -> str:
"""Search Wikipedia for a query and return maximum 2 page links.
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
# ]
# )
formatted_search_docs = [
{"source": doc.metadata.get("source", ""), "page": doc.metadata.get("page", "")}
for doc in search_docs
]
return {"wiki_page_links": json.dumps(formatted_search_docs)}
@tool
def web_search(query: str) -> str:
"""Search Tavily for a query and return maximum 3 page links.
Args:
query: The search query."""
search_docs = TavilySearchResults(
max_results=3,
tavily_api_key="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",
).invoke(query)
# formatted_search_docs = "\n\n---\n\n".join(
# [
# f'<Document source="{doc.get("url", "")}" title="{doc.get("title", "")}"/>\n{doc.get("content", "")}\n</Document>'
# for doc in search_docs
# ]
# )
formatted_search_docs = [
{"source": doc.get("url", ""), "page": doc.get("title", "")}
for doc in search_docs
]
return {"web_page_links": json.dumps(formatted_search_docs)}
@tool
def visit_webpage(url: str) -> str:
"""Retrieve content from a webpage using page links. A good option when you need detailed information from a specific webpage.
Args:
url: The URL to retrieve content from."""
search_docs = WebBaseLoader(url).load()
formatted_search_docs = "\n\n---\n\n".join(
[
f'<Document source="{doc.metadata.get("source", "")}" title="{doc.metadata.get("title", "")}"/>\n{doc.page_content}\n</Document>'
for doc in search_docs
]
)
return {"web_page_content": formatted_search_docs}
@tool
def arxiv_search(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["Title"]}" page="{doc.metadata.get("Summary", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
for doc in search_docs
]
)
return {"arxiv_results": formatted_search_docs}
@tool
def translate_to_english(text: str) -> str:
"""Translate the given text to English.
Args:
text: The text to translate to English."""
from langchain_ollama import ChatOllama
llm = ChatOllama(
model="tinyllama:1.1b",
api_base="http://localhost:11434",
)
prompt = f"Translate the following text to English:\n\n{text}"
response = llm.invoke(
[
{"role": "system", "content": "You are a helpful assistant that translates text to English."},
{"role": "user", "content": prompt},
]
)
return {"translated_text": response.content}
@tool
def multiply(a: float, b: float) -> float:
"""
Multiplies two numbers.
Args:
a (float): the first number
b (float): the second number
"""
return a * b
@tool
def add(a: float, b: float) -> float:
"""
Adds two numbers.
Args:
a (float): the first number
b (float): the second number
"""
return a + b
@tool
def subtract(a: float, b: float) -> int:
"""
Subtracts two numbers.
Args:
a (float): the first number
b (float): the second number
"""
return a - b
@tool
def divide(a: float, b: float) -> float:
"""
Divides two numbers.
Args:
a (float): the first float number
b (float): the second float number
"""
if b == 0:
raise ValueError("Cannot divided by zero.")
return a / b
@tool
def modulus(a: int, b: int) -> int:
"""
Get the modulus of two numbers.
Args:
a (int): the first number
b (int): the second number
"""
return a % b
@tool
def power(a: float, b: float) -> float:
"""
Get the power of two numbers.
Args:
a (float): the first number
b (float): the second number
"""
return a**b
@tool
def square_root(a: float) -> float | complex:
"""
Get the square root of a number.
Args:
a (float): the number to get the square root of
"""
if a >= 0:
return a**0.5
return cmath.sqrt(a) |