huytofu92's picture
enhance tool
26a59c8
from langchain_core.tools import tool as langchain_tool
from smolagents.tools import Tool, tool
from datetime import datetime
from typing import Literal, List, Union
from smolagents import VisitWebpageTool
from langchain_community.tools.tavily_search import TavilySearchResults
import pandas as pd
import os
@tool
def get_current_time(timezone: str = "America/New_York", format: str = "%Y-%m-%d %H:%M:%S")->str:
"""
Get the current time
Args:
timezone: The timezone to get the current time in. Example: "America/New_York"
format: The format to return the current time in. Example: "%Y-%m-%d %H:%M:%S"
Returns:
The current time
"""
return datetime.now(timezone).strftime(format)
@tool
def sort_list(my_list: List[int], order: Literal["asc", "desc", "alphabetize", "alphabetize_reverse"])->List[int]:
"""
Sort a list in ascending or descending order if the list contains numbers.
Sort it in alphabetically or alphabetically in reverse order if the list contains strings or mixed types.
Args:
my_list: The list to sort
order: The order to sort the list in. Must be one of the following:
- "asc": Sort the list in ascending order. Only for lists containing numbers.
- "desc": Sort the list in descending order. Only for lists containing numbers.
- "alphabetize": Sort the list alphabetically. Only for lists containing strings or mixed types.
- "alphabetize_reverse": Sort the list alphabetically in reverse order. Only for lists containing strings or mixed types.
Returns:
The sorted list
"""
if not isinstance(my_list, List):
raise ValueError("my_list must be a list")
else:
if all(isinstance(item, (int, float)) for item in my_list):
if order in ["asc", "desc"]:
return sorted(my_list, reverse=order == "desc")
elif order in ["alphabetize", "alphabetize_reverse"]:
how = {
"alphabetize": "asc",
"alphabetize_reverse": "desc"
}
return sorted(my_list, key=lambda x: str(x), reverse=how[order] == "desc")
else:
raise ValueError("order must be one of the following: asc, desc, alphabetize, alphabetize_reverse")
else:
print("This is a mixed list. Converting and sorting alphabetically.")
my_list = [str(item) for item in my_list]
how = {
"alphabetize": "asc",
"alphabetize_reverse": "desc"
}
return sorted(my_list, reverse=how[order] == "desc")
#smolagents tools
# visit_webpage_tool = VisitWebpageTool()
tavily_search_tool = Tool.from_langchain(TavilySearchResults(k=3))
@tool
def operate_two_numbers(num1: float, num2: float, operation: Literal["add", "subtract", "multiply", "divide", "power", "modulo"], decimal_places: int = 2)->float:
"""
Operate on two numbers
Args:
num1: The first number to operate on. Must be a float.
num2: The second number to operate on. Must be a float.
operation: The operation to perform. Must be one of the following:
- "add": Add the two numbers
- "subtract": Subtract the two numbers
- "multiply": Multiply the two numbers
- "divide": Divide the two numbers
- "power": Raise the first number to the power of the second number
- "modulo": Return the remainder of the division of the first number by the second number
decimal_places: The number of decimal places to round the result to. Default is 2.
Returns:
The result of the operation
"""
if operation == "add":
return round(num1 + num2, decimal_places)
elif operation == "subtract":
return round(num1 - num2, decimal_places)
elif operation == "multiply":
return round(num1 * num2, decimal_places)
elif operation == "divide":
return round(num1 / num2, decimal_places)
elif operation == "power":
return round(num1 ** num2, decimal_places)
elif operation == "modulo":
return round(num1 % num2, decimal_places)
else:
raise ValueError("operation must be one of the following: add, subtract, multiply, divide, power, modulo")
@tool
def convert_number(orig_num: any, operation: Literal["to_base", "type_cast"], new_base: Literal["binary", "octal", "hexadecimal", "int", "float"], decimal_places: int = 2)->any:
"""
Convert a number to a new base
Args:
orig_num: The number to convert. Must be a float or int.
operation: The operation to perform. Must be one of the following:
- "to_base": Convert the number to a new base.
- "type_cast": Convert the number to a new type.
new_base: The new base to convert the number to. Must be one of the following:
- "binary": Convert the number to binary.
- "octal": Convert the number to octal.
- "hexadecimal": Convert the number to hexadecimal.
- "int": Convert the number to an int.
- "float": Convert the number to a float.
decimal_places: The number of decimal places to round the result to. Default is 2. Only used if operation is "type_cast" and new_base is "float".
Returns:
The converted number. Can be float or int or str.
"""
if operation == "to_base":
if new_base == "binary":
return bin(orig_num)
elif new_base == "octal":
return oct(orig_num)
elif new_base == "hexadecimal":
return hex(orig_num)
else:
raise ValueError("new_base must be one of the following: binary, octal, hexadecimal, int, float")
elif operation == "type_cast":
if new_base == "int":
return int(orig_num)
elif new_base == "float":
return round(float(orig_num), decimal_places)
else:
raise ValueError("new_base must be one of the following: int, float")
else:
raise ValueError("operation must be one of the following: to_base, type_cast")
@tool
def load_dataframe_from_csv(file_path: str)->pd.DataFrame:
"""
Load a pandas DataFrame from a CSV file
Args:
file_path: The path to the CSV file to load.
Returns:
The pandas DataFrame
"""
return pd.read_csv(file_path)
@tool
def load_dataframe_from_excel(file_path: str)->pd.DataFrame:
"""
Load a pandas DataFrame from an Excel file
Args:
file_path: The path to the Excel file to load.
Returns:
The pandas DataFrame
"""
try:
df = pd.read_excel(file_path)
except Exception as e:
curr_dir = os.path.dirname(os.path.abspath(__file__))
file_path = os.path.join(curr_dir, file_path)
df = pd.read_excel(file_path)
return df
@tool
def to_dataframe(data: List[dict], columns: List[str])->pd.DataFrame:
"""
Convert a list of dictionaries to a pandas DataFrame
Args:
data: The list of dictionaries to convert to a pandas DataFrame.
columns: The columns of the pandas DataFrame.
Returns:
The pandas DataFrame
"""
return pd.DataFrame(data, columns=columns)
@tool
def to_json(data: pd.DataFrame)->str:
"""
Convert a pandas DataFrame to a JSON string
Args:
data: The pandas DataFrame to convert to a JSON string.
Returns:
The JSON string
"""
return data.to_json(orient="records")
@tool
def get_dataframe_data(data: pd.DataFrame, column: any, row: any)->any:
"""
Get a specific cell from a pandas DataFrame
Args:
data: The pandas DataFrame to get the data from.
column: The column to get the data from. Must be a string or int. If int then it is the index of the column.
row: The row to get the data from. Must be a string or int. If int then it is the index of the row.
Returns:
The data from the specified cell. Can be float or int or str.
"""
if isinstance(column, int):
column = data.iloc[:, column]
if isinstance(row, int):
row = data.iloc[row, :]
return data.loc[row, column]
@tool
def get_dataframe_column(data: pd.DataFrame, column: any)->pd.Series:
"""
Get a specific column from a pandas DataFrame
Args:
data: The pandas DataFrame to get the column from.
column: The column to get the data from. Must be a string or int. If int then it is the index of the column.
Returns:
The data from the specified column
"""
return data.iloc[:, column]
@tool
def get_dataframe_row(data: pd.DataFrame, row: any)->pd.Series:
"""
Get a specific row from a pandas DataFrame
Args:
data: The pandas DataFrame to get the row from.
row: The row to get the data from. Must be a string or int. If int then it is the index of the row.
Returns:
The data from the specified row
"""
return data.iloc[row, :]
@tool
def get_dataframe_groupby(data: pd.DataFrame, column: any, operation: Literal["mean", "sum", "count", "min", "max", "median", "std", "var"])->pd.DataFrame:
"""
Group a pandas DataFrame by a specific column and perform an operation on the grouped data
Args:
data: The pandas DataFrame to group.
column: The column to group the data by.
operation: The operation to perform on the grouped data. Must be one of the following:
- "mean": Calculate the mean of the grouped data.
- "sum": Calculate the sum of the grouped data.
- "count": Count the number of rows in the grouped data.
- "min": Calculate the minimum of the grouped data.
- "max": Calculate the maximum of the grouped data.
- "median": Calculate the median of the grouped data.
- "std": Calculate the standard deviation of the grouped data.
- "var": Calculate the variance of the grouped data.
Returns:
The grouped data
"""
if operation == "mean":
return data.groupby(column).mean()
elif operation == "sum":
return data.groupby(column).sum()
elif operation == "count":
return data.groupby(column).count()
elif operation == "min":
return data.groupby(column).min()
elif operation == "max":
return data.groupby(column).max()
elif operation == "median":
return data.groupby(column).median()
elif operation == "std":
return data.groupby(column).std()
elif operation == "var":
return data.groupby(column).var()
else:
raise ValueError("operation must be one of the following: mean, sum, count, min, max, median, std, var")
@tool
def read_python_file_from_path(file_path: str) -> str:
"""
Read and return the contents of a Python file from a given path.
Args:
file_path: Path to the Python file to read
Returns:
str: Contents of the Python file
"""
try:
# Check if file exists
# if not os.path.exists(file_path):
# raise FileNotFoundError(f"File not found: {file_path}")
# Check if it's a Python file
if not file_path.endswith('.py'):
raise ValueError(f"File is not a Python file: {file_path}")
# Try reading with absolute path first
try:
with open(file_path, 'r', encoding='utf-8') as f:
return f.read()
except Exception as e:
print(f"Failed to read with absolute path: {str(e)}")
# Try with adjusted path
current_file_path = os.path.abspath(__file__)
current_file_dir = os.path.dirname(current_file_path)
adjusted_path = os.path.join(current_file_dir, file_path)
print(f"Trying adjusted path: {adjusted_path}")
# if not os.path.exists(adjusted_path):
# raise FileNotFoundError(f"File not found at either {file_path} or {adjusted_path}")
with open(adjusted_path, 'r', encoding='utf-8') as f:
return f.read()
except Exception as e:
raise RuntimeError(f"Error reading Python file: {str(e)}")