add filtering tool

app.py

@@ -4,6 +4,7 @@ import requests
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
+import pandas as pd
 
 from Gradio_UI import GradioUI
 
@@ -33,6 +34,45 @@ def get_current_time_in_timezone(timezone: str) -> str:
     except Exception as e:
         return f"Error fetching time for timezone '{timezone}': {str(e)}"
 
+@tool
+def filter_dataframe(df: pd.DataFrame, filters: dict) -> pd.DataFrame:
+    """
+    Filters a DataFrame based on provided keys and their corresponding values.
+    
+    Parameters:
+        df (pd.DataFrame): The pandas DataFrame to filter.
+        filters (dict): A dictionary where each key is a column name in the DataFrame,
+                        and the corresponding value is the filter criteria. The filter 
+                        value can be a single value for an equality check, or a list of 
+                        values to use with the 'isin' filter.
+                        
+                        Example:
+                            filters = {
+                                "name": "Alice",          # filter where 'name' == "Alice"
+                                "city": ["New York", "Boston"]  # filter where 'city' is either "New York" or "Boston"
+                            }
+                            
+    Returns:
+        pd.DataFrame: A DataFrame filtered based on the provided criteria.
+        
+    Raises:
+        ValueError: If any of the provided filter keys are not present in the DataFrame.
+    """
+    filtered_df = df.copy()
+    
+    for column, value in filters.items():
+        if column not in filtered_df.columns:
+            raise ValueError(f"Column '{column}' does not exist in the DataFrame.")
+        
+        if isinstance(value, list):
+            # Filter for any value in the list (like SQL IN)
+            filtered_df = filtered_df[filtered_df[column].isin(value)]
+        else:
+            # Filter for equality
+            filtered_df = filtered_df[filtered_df[column] == value]
+    
+    return filtered_df
+
 
 final_answer = FinalAnswerTool()
 
@@ -55,7 +95,7 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer], ## add your tools here (don't remove final answer)
+        tools=[final_answer, get_current_time_in_timezone, filter_dataframe], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,