from smolagents import tool
import pandas as pd

@tool
def inspect_dataframe(file_path: str, operation: str, query_code: str = None) -> str:
    """
    Safely inspects or queries a CSV or Excel sheet using pandas.
    
    Args:
        file_path: The local path to the .csv or .xlsx file.
        operation: 'summary' to view columns/types/head, or 'query' to evaluate custom python string execution.
        query_code: The Python code snippet to run against the dataframe 'df' (e.g., "df['age'].mean()").
    """
    try:
        if file_path.endswith('.xlsx'):
            df = pd.read_excel(file_path)
        else:
            df = pd.read_csv(file_path)
            
        if operation == 'summary':
            summary = f"Columns: {list(df.columns)}\nShape: {df.shape}\nFirst 3 rows:\n{df.head(3).to_string()}"
            return summary
            
        elif operation == 'query' and query_code:
            # Safely evaluate code expecting a local variable 'df'
            result = eval(query_code, {"df": df, "pd": pd})
            return str(result)
            
        return "Invalid operation requested."
    except Exception as e:
        return f"Error processing file: {str(e)}"