|
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() |
|
|