redefine tools

tools.py

@@ -97,7 +97,7 @@ def operate_two_numbers(num1: float, num2: float, operation: Literal["add", "sub
         raise ValueError("operation must be one of the following: add, subtract, multiply, divide, power, modulo")
 
 @tool
-def convert_number(orig_num: Union[float, int], operation: Literal["to_base", "type_cast"], new_base: Literal["binary", "octal", "hexadecimal", "int", "float"], decimal_places: int = 2)->Union[int, float]:
+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:
@@ -113,7 +113,7 @@ def convert_number(orig_num: Union[float, int], operation: Literal["to_base", "t
             - "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
+        The converted number. Can be float or int or str.
     """
     if operation == "to_base":
         if new_base == "binary":
@@ -149,7 +149,7 @@ def to_json(data: pd.DataFrame)->str:
     return data.to_json(orient="records")
 
 @tool
-def get_dataframe_data(data: pd.DataFrame, column: Union[str, int], row: Union[str, int])->Union[str, int, float]:
+def get_dataframe_data(data: pd.DataFrame, column: any, row: any)->any:
     """
     Get a specific cell from a pandas DataFrame
     Args:
@@ -157,7 +157,7 @@ def get_dataframe_data(data: pd.DataFrame, column: Union[str, int], row: Union[s
         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
+        The data from the specified cell. Can be float or int or str.
     """
     if isinstance(column, int):
         column = data.iloc[:, column]
@@ -166,7 +166,7 @@ def get_dataframe_data(data: pd.DataFrame, column: Union[str, int], row: Union[s
     return data.loc[row, column]
 
 @tool
-def get_dataframe_column(data: pd.DataFrame, column: Union[str, int])->pd.Series:
+def get_dataframe_column(data: pd.DataFrame, column: any)->pd.Series:
     """
     Get a specific column from a pandas DataFrame
     Args:
@@ -178,7 +178,7 @@ def get_dataframe_column(data: pd.DataFrame, column: Union[str, int])->pd.Series
     return data.iloc[:, column]
 
 @tool
-def get_dataframe_row(data: pd.DataFrame, row: Union[str, int])->pd.Series:
+def get_dataframe_row(data: pd.DataFrame, row: any)->pd.Series:
     """
     Get a specific row from a pandas DataFrame
     Args:
@@ -190,7 +190,7 @@ def get_dataframe_row(data: pd.DataFrame, row: Union[str, int])->pd.Series:
     return data.iloc[row, :]
 
 @tool 
-def get_dataframe_groupby(data: pd.DataFrame, column: str, operation: Literal["mean", "sum", "count", "min", "max", "median", "std", "var"])->pd.DataFrame:
+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: