mchinea
update tools and model
afadec7
Raw
History Blame Contribute Delete
8.4 kB
import os
import random
import requests
import tempfile
import re
from typing import Dict
from pathlib import Path
#from markitdown import MarkItDown
from urllib.parse import urlparse
from langchain_core.tools import tool
from langchain_core.messages import ToolMessage
from langchain_tavily import TavilySearch
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.document_loaders import ArxivLoader
@tool
def web_search(query: str) -> ToolMessage:
"""Search in the web with Tavily for a query and return maximum 5 results.
Args:
query: The search query.
Returns:
Tavily output, and snippet for the top 5 results
"""
return TavilySearch(max_results=5, include_images=False).invoke({"query": query})
@tool
def search_tool(query: str) -> str:
"""Search in Google and returns an string with title, link, and snippet for the top 5 results.
Args:
query: str
Returns:
Title, link, and snippet for the top 5 results
"""
searcher = GoogleSerperAPIWrapper(k=5)
retries = 3
result = ""
while retries > 0:
try:
search_results = searcher.results(query)["organic"]
for row in search_results:
result += f"Title: {row['title']}\nSnippet: {row['snippet']}\nURL: {row['link']}\n\n"
return result
except Exception as e:
retries -= 1
return f"There was an error with Google search: {e}"
@tool
def wikipedia_search(query: str) -> Dict[str, list]:
"""Search Wikipedia for a given query and return the first 10 results.
Args:
query: The search term or topic.
Returns:
A dictionary containing the formatted Wikipedia results.
"""
search_docs = WikipediaLoader(query=query, load_max_docs=10).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 {"wiki_results": formatted_search_docs}
#Mathematical tools
@tool
def multiply(a: float, b: float) -> float:
"""Multiply two numbers.
Args:
a: first number
b: second number
Returns:
Multiplication result
"""
return a * b
@tool
def add(a: float, b: float) -> float:
"""Add two numbers.
Args:
a: first number
b: second number
Returns:
Addition result
"""
return a + b
@tool
def subtract(a: float, b: float) -> float:
"""Subtract two numbers.
Args:
a: first number
b: second number
Returns:
Subtraction result
"""
return a - b
@tool
def divide(a: float, b: float) -> float:
"""Divide two numbers.
Args:
a: first number
b: second number
Returns:
Division result
"""
if b == 0:
raise ValueError("Cannot divide by zero.")
return a / b
@tool
def modulus(a: int, b: int) -> int:
"""Get the modulus of two numbers.
Args:
a: first number
b: second number
Returns:
Modulus result
"""
return a % b
from langchain_core.tools import tool
@tool
def convert_units(value: float, from_unit: str, to_unit: str) -> float:
"""
Converts a value from one unit to another.
Args:
value: The numerical value to convert.
from_unit: The original unit (e.g. 'miles', 'kg', 'celsius').
to_unit: The target unit (e.g. 'kilometers', 'lb', 'fahrenheit').
Supported conversions:
- miles <-> kilometers
- kilograms <-> pounds
- celsius <-> fahrenheit
Returns:
The converted value result.
"""
conversions = {
("miles", "kilometers"): lambda v: v * 1.60934,
("kilometers", "miles"): lambda v: v / 1.60934,
("kilograms", "pounds"): lambda v: v * 2.20462,
("pounds", "kilograms"): lambda v: v / 2.20462,
("celsius", "fahrenheit"): lambda v: (v * 9/5) + 32,
("fahrenheit", "celsius"): lambda v: (v - 32) * 5/9,
}
key = (from_unit.lower(), to_unit.lower())
if key not in conversions:
raise ValueError(f"Conversion from {from_unit} to {to_unit} not supported.")
return conversions[key](value)
@tool
def query_table_data(file_path: str, query: str, sheet_name: str = None) -> str:
"""
Loads a table from CSV or Excel and filters it using a pandas query.
Args:
file_path: Path to the table file (.xlsx, .xls).
query_pandas_syntax: A pandas-compatible query string, e.g., "Age > 30 and Country == 'USA'".
sheet_name: Optional sheet name if the file is Excel.
Returns:
A string representation (markdown) of the filtered table (max 10 rows).
"""
try:
import pandas as pd
path = Path(file_path)
if not path.exists():
raise FileNotFoundError(f"File not found: {file_path}")
ext = path.suffix.lower()
if ext == ".csv":
df = pd.read_csv(path)
elif ext in [".xlsx", ".xls"]:
df = pd.read_excel(path, sheet_name=sheet_name)
else:
raise ValueError(f"Unsupported file extension: {ext}")
try:
#Converts a natural language query to pandas query syntax using basic heuristics.
# Preprocess query
query_l = query.lower().strip()
# Heuristic rules
rules = [
(r"(\w+) greater than (\d+)", r"\1 > \2"),
(r"(\w+) less than (\d+)", r"\1 < \2"),
(r"(\w+) equal to ['\"]?([\w\s]+)['\"]?", r"\1 == '\2'"),
(r"(\w+) not equal to ['\"]?([\w\s]+)['\"]?", r"\1 != '\2'"),
(r"(\w+) more than (\d+)", r"\1 > \2"),
(r"(\w+) less than or equal to (\d+)", r"\1 <= \2"),
(r"(\w+) greater than or equal to (\d+)", r"\1 >= \2"),
(r"(\w+) is ['\"]?([\w\s]+)['\"]?", r"\1 == '\2'"),
]
for pattern, replacement in rules:
if re.search(pattern, query):
query = re.sub(pattern, replacement, query)
break
# Handle AND/OR logic
query_pandas_syntax = query.replace(" and ", " and ")
query_pandas_syntaxs = query.replace(" or ", " or ")
filtered_df = df.query(query_pandas_syntax)
return filtered_df.head(10).to_markdown(index=False)
except Exception as e:
raise ValueError(f"Invalid query: {query_pandas_syntax}. Error: {e}")
except ImportError:
return "Error: pandas and openpyxl are not installed. Please install them with 'pip install pandas openpyxl'."
@tool
def arvix_search(query: str) -> str:
"""Search Arxiv for a query and return maximum 5 result.
Args:
query: The search query.
Returns:
A dictionary containing the formatted Arvix results, and snippet for the top 5 results.
"""
search_docs = ArxivLoader(query=query, load_max_docs=5).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 {"arvix_results": formatted_search_docs}
@tool
def read_python_file(file_path: str) -> str:
"""
Reads and parses an Python file to markdown.
Args:
file_path: Path to the Python file
Returns:
Python file content.
"""
try:
# Just with markitdown
path = Path(file_path)
if not path.exists():
raise FileNotFoundError(f"File not found: {file_path}")
ext = path.suffix.lower()
if ext == ".py":
md = MarkItDown(enable_plugins=True)
result = md.convert(file_path)
return result.text_content
else:
raise ValueError(f"Unsupported file extension: {ext}")
except Exception as err:
raise type(err)(f"Could not parse python file > {err}")
level1_tools = [
multiply,
add,
subtract,
divide,
modulus,
wikipedia_search,
web_search,
#search_tool,
arvix_search,
convert_units,
query_table_data,
read_python_file,
]