/ examples / tools / example_pandas_tool.py
example_pandas_tool.py
  1  """
  2  Example: Using pandas as a custom tool with PraisonAI Agent
  3  
  4  This example shows how to wrap pandas functions as tools
  5  that can be used with Agent().
  6  """
  7  
  8  import pandas as pd
  9  from io import StringIO
 10  
 11  # Define pandas wrapper tools
 12  def analyze_csv_data(csv_data: str) -> str:
 13      """Analyze CSV data using pandas DataFrame.
 14      
 15      Args:
 16          csv_data: CSV-formatted string data
 17          
 18      Returns:
 19          Analysis summary including shape, columns, and statistics
 20      """
 21      df = pd.read_csv(StringIO(csv_data))
 22      summary = f"""DataFrame Analysis:
 23  - Shape: {df.shape[0]} rows, {df.shape[1]} columns
 24  - Columns: {list(df.columns)}
 25  - Data Types: {dict(df.dtypes)}
 26  - Summary Statistics:
 27  {df.describe().to_string()}"""
 28      return summary
 29  
 30  
 31  def filter_dataframe(csv_data: str, column: str, condition: str, value: str) -> str:
 32      """Filter a DataFrame based on a condition.
 33      
 34      Args:
 35          csv_data: CSV-formatted string data
 36          column: Column name to filter on
 37          condition: Condition type ('equals', 'greater', 'less', 'contains')
 38          value: Value to compare against
 39          
 40      Returns:
 41          Filtered data as CSV string
 42      """
 43      df = pd.read_csv(StringIO(csv_data))
 44      
 45      if condition == 'equals':
 46          filtered = df[df[column] == value]
 47      elif condition == 'greater':
 48          filtered = df[df[column] > float(value)]
 49      elif condition == 'less':
 50          filtered = df[df[column] < float(value)]
 51      elif condition == 'contains':
 52          filtered = df[df[column].str.contains(value, na=False)]
 53      else:
 54          return f"Unknown condition: {condition}"
 55      
 56      return filtered.to_csv(index=False)
 57  
 58  
 59  def calculate_statistics(csv_data: str, column: str) -> str:
 60      """Calculate statistics for a specific column.
 61      
 62      Args:
 63          csv_data: CSV-formatted string data
 64          column: Column name to analyze
 65          
 66      Returns:
 67          Statistics summary for the column
 68      """
 69      df = pd.read_csv(StringIO(csv_data))
 70      
 71      if column not in df.columns:
 72          return f"Column '{column}' not found. Available: {list(df.columns)}"
 73      
 74      col = df[column]
 75      
 76      if pd.api.types.is_numeric_dtype(col):
 77          stats = f"""Statistics for '{column}':
 78  - Count: {col.count()}
 79  - Mean: {col.mean():.2f}
 80  - Std: {col.std():.2f}
 81  - Min: {col.min()}
 82  - Max: {col.max()}
 83  - Median: {col.median()}"""
 84      else:
 85          stats = f"""Statistics for '{column}':
 86  - Count: {col.count()}
 87  - Unique: {col.nunique()}
 88  - Top: {col.mode().iloc[0] if len(col.mode()) > 0 else 'N/A'}
 89  - Sample values: {list(col.head(3))}"""
 90      
 91      return stats
 92  
 93  
 94  if __name__ == "__main__":
 95      # Test data
 96      test_csv = """name,age,salary,department
 97  Alice,30,50000,Engineering
 98  Bob,25,45000,Marketing
 99  Charlie,35,60000,Engineering
100  Diana,28,52000,Sales
101  Eve,32,55000,Marketing"""
102  
103      print("=== Test: analyze_csv_data ===")
104      print(analyze_csv_data(test_csv))
105      print()
106      
107      print("=== Test: filter_dataframe ===")
108      print(filter_dataframe(test_csv, "age", "greater", "28"))
109      print()
110      
111      print("=== Test: calculate_statistics ===")
112      print(calculate_statistics(test_csv, "salary"))
113      print()
114      
115      # Example with Agent (requires praisonaiagents)
116      print("=== Agent Usage Example ===")
117      print("""
118  from praisonaiagents import Agent
119  
120  agent = Agent(
121      name="data_analyst",
122      role="Data Analyst",
123      goal="Analyze data using pandas tools",
124      tools=[analyze_csv_data, filter_dataframe, calculate_statistics]
125  )
126  
127  result = agent.start("Analyze this data and find the average salary")
128  """)