File size: 2,837 Bytes
30c51ff |
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 |
import pytesseract
from PIL import Image
from dotenv import load_dotenv
from langchain.tools import tool
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.tools import TavilySearchResults
@tool
def multiply(a: int, b: int) -> int:
"""Multiply two numbers.
Args:
a: first int
b: second int
"""
return a * b
@tool
def add(a: int, b: int) -> int:
"""Add two numbers.
Args:
a: first int
b: second int
"""
return a + b
@tool
def subtract(a: int, b: int) -> int:
"""Subtract two numbers.
Args:
a: first int
b: second int
"""
return a - b
@tool
def divide(a: int, b: int) -> float:
"""Divide two numbers.
Args:
a: first int
b: second int
"""
try:
return a / b
except ZeroDivisionError:
raise ValueError("Cannot divide by zero.")
@tool
def modulus(a: int, b: int) -> int:
"""Get the modulus of two numbers.
Args:
a: first int
b: second int
"""
return a % b
@tool
def wiki_search(query: str) -> str:
"""Search Wikipedia for a query and return up to 3 results.
Args:
query: query to search.
"""
docs = WikipediaLoader(query, load_max_docs=3).load()
results = "wiki_search results:\n\n"
results += '\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 docs
])
return results
@tool
def web_search(query: str) -> str:
"""Search the web for a query and return up to 3 results.
Args:
query: query to search.
"""
docs = TavilySearchResults(max_results=3).invoke(query)
results = "web_search results:\n\n"
results += '\n\n---\n\n'.join([
f'<Document source="{doc["url"]}"/>\n{doc["content"]}\n</Document>'
for doc in docs
])
return results
@tool
def extract_text_from_image(image_path: str) -> str:
"""Extract text from a image.
Args:
image_path: path to the image
Returns:
extracted text from the image
"""
try:
image = Image.open(image_path)
text = pytesseract.image_to_string(image)
return f'Extracted text: {text}'
except Exception as err:
return f'Error extracting text from the image {image_path}: {str(err)}'
all_tools = [
multiply,
add,
subtract,
divide,
modulus,
wiki_search,
web_search,
extract_text_from_image,
]
def main():
load_dotenv()
results = wiki_search('What is Dijkstra algorithm?')
print(results)
print('-'*80)
results = web_search('What is TypedDict in python')
print(results)
print('-'*80)
if __name__ == '__main__':
main()
|