Инструменты агентов на основе Pandas

Инструменты pandas для агентов данных

Инструмент: read_csv

import pandas as pd
from typing import Optional

# In-memory dataframe store (shared across tool calls in one session)
dataframes = {}

def read_csv(path: str, df_name: str = 'df') -> dict:
    '''Load a CSV file and return schema info.'''
    df = pd.read_csv(path)
    dataframes[df_name] = df
    return {
        'df_name': df_name,
        'rows': len(df),
        'columns': list(df.columns),
        'dtypes': df.dtypes.astype(str).to_dict(),
        'sample': df.head(3).to_dict(orient='records'),
        'null_counts': df.isnull().sum().to_dict()
    }

# Agent usage:
result = read_csv('orders.csv', 'orders')
print(f"Loaded {result['rows']} rows, columns: {result['columns']}")

Инструмент: filter_dataframe

def filter_dataframe(df_name: str, condition: str, result_name: str = None) -> dict:
    '''Filter rows matching condition. Returns filtered dataframe info.'''
    df = dataframes.get(df_name)
    if df is None:
        return {'error': f'Dataframe {df_name!r} not found. Call read_csv first.'}

    try:
        filtered = df.query(condition)
        store_name = result_name or f'{df_name}_filtered'
        dataframes[store_name] = filtered
        return {
            'df_name': store_name,
            'rows_before': len(df),
            'rows_after': len(filtered),
            'sample': filtered.head(3).to_dict(orient='records')
        }
    except Exception as e:
        return {'error': f'Filter error: {e}'}

# Example: filter orders from 2024
result = filter_dataframe('orders', 'order_date >= "2024-01-01"', 'orders_2024')
print(f"Filtered to {result['rows_after']} rows")

Инструмент: groupby_aggregate

def groupby_aggregate(df_name: str, group_column: str,
                      agg_column: str, agg_func: str) -> dict:
    '''Group by column and aggregate. agg_func: sum, mean, count, min, max, median'''
    df = dataframes.get(df_name)
    if df is None:
        return {'error': f'Dataframe {df_name!r} not found'}

    VALID_FUNCS = {'sum', 'mean', 'count', 'min', 'max', 'median', 'std'}
    if agg_func not in VALID_FUNCS:
        return {'error': f'Invalid agg_func. Choose from: {VALID_FUNCS}'}

    try:
        result = df.groupby(group_column)[agg_column].agg(agg_func).reset_index()
        result.columns = [group_column, f'{agg_column}_{agg_func}']
        result = result.sort_values(f'{agg_column}_{agg_func}', ascending=False)
        return {
            'result': result.head(20).to_dict(orient='records'),
            'total_groups': len(result)
        }
    except Exception as e:
        return {'error': str(e)}

# Agent call: total revenue by product category
print(groupby_aggregate('orders', 'category', 'revenue', 'sum'))

Инструмент: merge_dataframes

def merge_dataframes(df1_name: str, df2_name: str,
                     on: str, how: str = 'inner',
                     result_name: str = 'merged') -> dict:
    '''Join two dataframes on a key column.'''
    df1 = dataframes.get(df1_name)
    df2 = dataframes.get(df2_name)
    if df1 is None:
        return {'error': f'{df1_name!r} not found'}
    if df2 is None:
        return {'error': f'{df2_name!r} not found'}

    VALID_HOW = {'inner', 'left', 'right', 'outer'}
    if how not in VALID_HOW:
        return {'error': f'Invalid join type. Choose from: {VALID_HOW}'}

    try:
        merged = df1.merge(df2, on=on, how=how)
        dataframes[result_name] = merged
        return {
            'df_name': result_name,
            'rows': len(merged),
            'columns': list(merged.columns),
            'sample': merged.head(3).to_dict(orient='records')
        }
    except Exception as e:
        return {'error': str(e)}

# Join orders with customers on customer_id
print(merge_dataframes('orders', 'customers', on='customer_id'))

Инструмент: describe_dataframe

def describe_dataframe(df_name: str) -> dict:
    '''Return statistical summary and data quality info.'''
    df = dataframes.get(df_name)
    if df is None:
        return {'error': f'Dataframe {df_name!r} not found'}

    # Numeric stats
    numeric_cols = df.select_dtypes(include='number').columns.tolist()
    stats = {}
    for col in numeric_cols:
        s = df[col].describe()
        stats[col] = {
            'min':    round(float(s['min']), 4),
            'max':    round(float(s['max']), 4),
            'mean':   round(float(s['mean']), 4),
            'median': round(float(df[col].median()), 4),
            'std':    round(float(s['std']), 4),
            'nulls':  int(df[col].isnull().sum())
        }

    # Categorical columns summary
    cat_cols = df.select_dtypes(include='object').columns.tolist()
    cat_stats = {
        col: {
            'unique_values': df[col].nunique(),
            'top_3': df[col].value_counts().head(3).to_dict()
        }
        for col in cat_cols
    }

    return {'numeric': stats, 'categorical': cat_stats, 'total_rows': len(df)}

Инструмент: sort_and_top_n

def sort_and_top_n(df_name: str, sort_column: str,
                   n: int = 10, ascending: bool = False) -> dict:
    '''Return top N rows sorted by column.'''
    df = dataframes.get(df_name)
    if df is None:
        return {'error': f'{df_name!r} not found'}
    if sort_column not in df.columns:
        return {'error': f'Column {sort_column!r} not found. Available: {list(df.columns)}'}

    sorted_df = df.sort_values(sort_column, ascending=ascending)
    top = sorted_df.head(n)

    return {
        'rows': top.to_dict(orient='records'),
        'total_rows_in_df': len(df),
        'sort_column': sort_column,
        'sort_order': 'ascending' if ascending else 'descending'
    }

# Top 5 orders by total value
print(sort_and_top_n('orders', 'total_amount', n=5))

Инструмент: pivot_table

def pivot_table(df_name: str, index: str, columns: str,
                values: str, aggfunc: str = 'sum') -> dict:
    '''Create a pivot table.'''
    df = dataframes.get(df_name)
    if df is None:
        return {'error': f'{df_name!r} not found'}

    VALID_FUNCS = {'sum', 'mean', 'count', 'min', 'max'}
    if aggfunc not in VALID_FUNCS:
        return {'error': f'Invalid aggfunc. Choose from: {VALID_FUNCS}'}

    try:
        pivot = df.pivot_table(
            index=index, columns=columns,
            values=values, aggfunc=aggfunc,
            fill_value=0
        )
        # Convert to nested dict for JSON serialization
        result = pivot.round(2).to_dict()
        return {
            'pivot': result,
            'index': index,
            'columns': columns,
            'values': values,
            'aggfunc': aggfunc
        }
    except Exception as e:
        return {'error': str(e)}

Инструмент: calculate_metric

def calculate_metric(df_name: str, metric: str, params: dict = None) -> dict:
    '''Calculate a named business metric from the dataframe.'''
    df = dataframes.get(df_name)
    if df is None:
        return {'error': f'{df_name!r} not found'}

    params = params or {}
    try:
        if metric == 'conversion_rate':
            total_col = params.get('total_column', 'total')
            converted_col = params.get('converted_column', 'converted')
            rate = df[converted_col].sum() / df[total_col].sum() * 100
            return {'metric': 'conversion_rate', 'value': round(rate, 2), 'unit': '%'}

        elif metric == 'average_order_value':
            revenue_col = params.get('revenue_column', 'revenue')
            count_col = params.get('count_column', 'order_count')
            aov = df[revenue_col].sum() / df[count_col].sum()
            return {'metric': 'average_order_value', 'value': round(aov, 2)}

        else:
            return {'error': f'Unknown metric: {metric}'}

    except Exception as e:
        return {'error': str(e)}

Регистрация инструментов в агенте на базе LLM

tools = [
    {
        'type': 'function',
        'function': {
            'name': 'read_csv',
            'description': 'Load a CSV file for analysis',
            'parameters': {
                'type': 'object',
                'properties': {
                    'path': {'type': 'string', 'description': 'Path to CSV file'},
                    'df_name': {'type': 'string', 'description': 'Name to reference this dataframe'}
                },
                'required': ['path']
            }
        }
    },
    {
        'type': 'function',
        'function': {
            'name': 'groupby_aggregate',
            'description': 'Group data by a column and compute sum, mean, count, min, max, or median',
            'parameters': {
                'type': 'object',
                'properties': {
                    'df_name':      {'type': 'string'},
                    'group_column': {'type': 'string'},
                    'agg_column':   {'type': 'string'},
                    'agg_func':     {'type': 'string', 'enum': ['sum', 'mean', 'count', 'min', 'max', 'median']}
                },
                'required': ['df_name', 'group_column', 'agg_column', 'agg_func']
            }
        }
    }
]

if __name__ == '__main__':
    print('Registered tools:')
    for t in tools:
        fn = t['function']
        print(f"  - {fn['name']}: {fn['description']}")

Объединение нескольких инструментов в цепочку

import json
import openai
import os

client = openai.OpenAI(api_key=os.getenv('OPENAI_API_KEY'))

TOOL_MAP = {
    'read_csv': read_csv,
    'filter_dataframe': filter_dataframe,
    'groupby_aggregate': groupby_aggregate,
    'describe_dataframe': describe_dataframe,
    'merge_dataframes': merge_dataframes,
    'sort_and_top_n': sort_and_top_n
}

def run_data_agent(question, data_paths):
    messages = [
        {'role': 'system', 'content': f'You are a data analyst. Available data: {data_paths}'},
        {'role': 'user', 'content': question}
    ]

    for _ in range(10):  # max 10 tool calls
        response = client.chat.completions.create(
            model='gpt-4o', messages=messages, tools=tools
        )
        msg = response.choices[0].message

        if not msg.tool_calls:
            return msg.content  # final answer

        messages.append(msg)
        for call in msg.tool_calls:
            fn = TOOL_MAP[call.function.name]
            args = json.loads(call.function.arguments)
            result = fn(**args)
            messages.append({'role': 'tool', 'tool_call_id': call.id,
                             'content': json.dumps(result)})

    return 'Max tool calls reached'

Проверка знаний

Повторение: инструменты агента данных на основе pandas