Excel Summarizer Agent Example

# schemas.py
from pydantic import BaseModel, Field, FilePath
from typing import List, Dict, Any, Optional, Union
from karo.schemas.base_schemas import BaseInputSchema, BaseOutputSchema
from karo.tools.base_tool import BaseToolInputSchema, BaseToolOutputSchema

# Schema for the Excel Reader Tool Input
class ExcelReaderInput(BaseToolInputSchema):
    """Input schema for the ExcelReaderTool."""
    file_path: FilePath = Field(..., description="The path to the Excel file (.xlsx or .xls).")
    sheet_name: Optional[Union[str, int]] = Field(None, description="Specific sheet name or index (0-based) to read. Reads the first sheet if None.")
    max_rows: Optional[int] = Field(100, description="Maximum number of rows to read from the sheet to avoid overly long context.")
    max_cols: Optional[int] = Field(20, description="Maximum number of columns to read.")

# Schema for the Excel Reader Tool Output
class ExcelReaderOutput(BaseToolOutputSchema):
    """Output schema for the ExcelReaderTool."""
    file_path: str = Field(..., description="The path of the file that was read.")
    sheet_name_read: str = Field(..., description="The name of the sheet that was actually read.")
    data_preview: Optional[str] = Field(None, description="A string representation (e.g., markdown table or CSV) of the first few rows/columns of the data.")
    row_count: Optional[int] = Field(None, description="Total number of rows read (up to max_rows).")
    column_names: Optional[List[str]] = Field(None, description="List of column names read (up to max_cols).")
    # Inherits success and error_message

# Schema for the Agent's final summary output
class ExcelSummaryOutput(BaseOutputSchema):
    """Output schema for the Excel Summarizer Agent."""
    summary: str = Field(..., description="A concise summary of the key insights from the Excel data.")
    key_findings: List[str] = Field(default_factory=list, description="A list of bullet points highlighting key findings.")
    potential_issues: Optional[str] = Field(None, description="Any potential issues or anomalies noted in the data.")

# Agent input remains simple for this example
class ExcelSummarizerInput(BaseInputSchema):
    """Input schema for the Excel Summarizer Agent."""
    file_path: str = Field(..., description="Path to the Excel file to be summarized.")
    user_query: Optional[str] = Field("Summarize the provided Excel data.", description="Optional user query guiding the summary.")

# tools/excel_reader_tool.py
import pandas as pd
from typing import Optional, List, Any, Union, Type, Dict
from pydantic import Field, FilePath
import logging
import os

# Import Karo base tool components using absolute path (assuming karo is installed)
from karo.tools.base_tool import BaseTool, BaseToolInputSchema, BaseToolOutputSchema
# Import schemas using relative path (assuming schemas.py is in the parent directory)
from ..schemas import ExcelReaderInput, ExcelReaderOutput

logger = logging.getLogger(__name__)

class ExcelReaderTool(BaseTool):
    """
    Reads data from a specified Excel file and sheet, providing a preview.
    Requires 'pandas' and 'openpyxl' to be installed.
    """
    input_schema: Type[ExcelReaderInput] = ExcelReaderInput
    output_schema: Type[ExcelReaderOutput] = ExcelReaderOutput
    name: str = "excel_file_reader"
    description: str = "Reads data from a specified sheet in an Excel (.xlsx, .xls) file and returns a preview string."

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """Initialize the ExcelReaderTool."""
        logger.info("ExcelReaderTool initialized.")
        # Check for pandas dependency during initialization
        try:
            import pandas
        except ImportError:
            logger.error("Missing dependency: 'pandas' is required for ExcelReaderTool.")
            raise ImportError("Missing dependency: 'pandas' is required for ExcelReaderTool. Install with `pip install pandas`.")
        # Check for openpyxl dependency during initialization
        try:
            import openpyxl
        except ImportError:
            logger.error("Missing dependency: 'openpyxl' is required for ExcelReaderTool.")
            raise ImportError("Missing dependency: 'openpyxl' is required for ExcelReaderTool. Install with `pip install openpyxl`.")
        pass # No specific config needed beyond dependency checks

    def run(self, input_data: ExcelReaderInput) -> ExcelReaderOutput:
        """
        Reads the specified Excel file and returns a data preview.

        Args:
            input_data: An instance of ExcelReaderInput.

        Returns:
            An instance of ExcelReaderOutput.
        """
        if not isinstance(input_data, self.input_schema):
             logger.warning(f"Invalid input type received: {type(input_data)}")
             # Attempt to create from dict if possible, otherwise fail
             try:
                 input_data = self.input_schema(**input_data)
             except Exception:
                 return self.output_schema(success=False, error_message="Invalid input data format.", file_path="N/A", sheet_name_read="N/A")


        file_path_str = str(input_data.file_path) # Handle potential Path object

        if not os.path.exists(file_path_str):
             logger.error(f"File not found: {file_path_str}")
             return self.output_schema(success=False, error_message=f"File not found at path: {file_path_str}", file_path=file_path_str, sheet_name_read="N/A")

        logger.info(f"Reading Excel file: {file_path_str}, Sheet: {input_data.sheet_name}, Max Rows: {input_data.max_rows}, Max Cols: {input_data.max_cols}")

        try:
            # Determine sheet to read
            excel_file = pd.ExcelFile(file_path_str, engine='openpyxl')
            sheet_names = excel_file.sheet_names
            sheet_to_read: Union[str, int] = 0
            sheet_name_read: str = sheet_names[0]

            if input_data.sheet_name is not None:
                if isinstance(input_data.sheet_name, int):
                    if 0 <= input_data.sheet_name < len(sheet_names):
                        sheet_to_read = input_data.sheet_name
                        sheet_name_read = sheet_names[sheet_to_read]
                    else:
                        err_msg = f"Sheet index {input_data.sheet_name} out of range (0-{len(sheet_names)-1})."
                        logger.warning(err_msg)
                        return self.output_schema(success=False, error_message=err_msg, file_path=file_path_str, sheet_name_read="N/A")
                elif isinstance(input_data.sheet_name, str):
                    if input_data.sheet_name in sheet_names:
                        sheet_to_read = input_data.sheet_name
                        sheet_name_read = input_data.sheet_name
                    else:
                        err_msg = f"Sheet name '{input_data.sheet_name}' not found. Available: {sheet_names}"
                        logger.warning(err_msg)
                        return self.output_schema(success=False, error_message=err_msg, file_path=file_path_str, sheet_name_read="N/A")

            # Read header to determine columns to use
            header_df = pd.read_excel(excel_file, sheet_name=sheet_to_read, nrows=0)
            all_columns = header_df.columns.tolist()
            cols_to_use_indices = list(range(min(len(all_columns), input_data.max_cols))) if input_data.max_cols else None
            cols_to_use_names = [all_columns[i] for i in cols_to_use_indices] if cols_to_use_indices is not None else None


            # Read data with limits
            df = pd.read_excel(
                excel_file,
                sheet_name=sheet_to_read,
                nrows=input_data.max_rows,
                usecols=cols_to_use_names # Use names derived from indices
            )

            # Generate preview string (e.g., markdown table)
            preview_rows = min(len(df), 10) # Limit preview rows
            data_preview_str = df.head(preview_rows).to_markdown(index=False)

            logger.info(f"Successfully read {len(df)} rows and {len(df.columns)} columns from sheet '{sheet_name_read}'.")

            return self.output_schema(
                success=True,
                file_path=file_path_str,
                sheet_name_read=sheet_name_read,
                data_preview=data_preview_str,
                row_count=len(df),
                column_names=df.columns.tolist() # Return actual columns read
            )

        except Exception as e:
            logger.error(f"Failed to read Excel file '{file_path_str}': {e}", exc_info=True)
            return self.output_schema(success=False, error_message=f"Error reading Excel file: {e}", file_path=file_path_str, sheet_name_read="N/A")

# excel_summarizer_agent.py
import argparse
import os
from dotenv import load_dotenv
from rich.console import Console
from rich.panel import Panel
import logging

# Configure basic logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Load environment variables
load_dotenv()

# Import Karo components and our custom items
from karo.core.base_agent import BaseAgent, BaseAgentConfig
from karo.providers.openai_provider import OpenAIProvider, OpenAIProviderConfig
from karo.prompts.system_prompt_builder import SystemPromptBuilder
from karo.schemas.base_schemas import AgentErrorSchema
# Import custom schemas and tool
from schemas import ExcelSummarizerInput, ExcelSummaryOutput # Assumes schemas.py is in the same dir
from tools.excel_reader_tool import ExcelReaderTool # Assumes tools/excel_reader_tool.py

# Initialize console
console = Console()

def run_summarization(file_path: str):
    """Orchestrates reading the Excel file and generating a summary."""
    console.print(Panel(f"[bold magenta]Karo Framework - Excel Summarizer Agent[/bold magenta]", title="Welcome", expand=False))
    console.print(f"Processing file: [cyan]{file_path}[/cyan]")

    # 1. Check API Key
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        console.print("[bold red]Error:[/bold red] OPENAI_API_KEY needed.")
        return

    # 2. Initialize Tool
    try:
        excel_reader = ExcelReaderTool()
        available_tools = [excel_reader]
        console.print("[green]✓ Excel Reader Tool Initialized[/green]")
    except ImportError as e: # Catch missing dependencies
         console.print(f"[bold red]Error initializing tools: {e}")
         console.print("  Please install required dependencies: pip install pandas openpyxl")
         return
    except Exception as e:
        console.print(f"[bold red]Error initializing tools: {e}")
        return

    # 3. Initialize Provider
    try:
        provider_config = OpenAIProviderConfig(model="gpt-4o-mini") # Or GPT-4 for potentially better summarization
        provider = OpenAIProvider(config=provider_config)
        console.print(f"[green]✓ OpenAI Provider Initialized (Model: {provider.get_model_name()})[/green]")
    except Exception as e:
        console.print(f"[bold red]Error initializing provider: {e}")
        return

    # 4. Configure Prompt Builder
    # This prompt guides the agent AFTER the tool has run and provided data.
    # The BaseAgent currently doesn't automatically inject tool output into a *second* LLM call.
    # So, this example focuses on the tool execution part first.
    # A more advanced agent would handle the two-step process (read -> summarize).
    # For now, we'll configure the agent to use the tool, but the summarization
    # part would need to be handled differently (e.g., in the calling script
    # after getting the tool output, or by extending BaseAgent).

    system_prompt_for_tool_use = """
You are an assistant tasked with reading data from Excel files.
When the user provides a file path, use the 'excel_file_reader' tool to read it.
Return the preview data provided by the tool.
"""
    prompt_builder = SystemPromptBuilder(
        role_description="Excel Data Reader Assistant",
        core_instructions=system_prompt_for_tool_use
    )

    # 5. Configure Agent
    # Note: We use the tool's output schema here just to demonstrate tool integration.
    # A real summarizer would likely use a different agent output schema.
    agent_config = BaseAgentConfig(
        provider=provider,
        tools=available_tools,
        prompt_builder=prompt_builder,
        input_schema=ExcelSummarizerInput, # Agent takes file path
        output_schema=ExcelReaderOutput # Agent's goal is to return tool output for now
    )
    agent = BaseAgent(config=agent_config)
    console.print("[green]✓ Excel Reader Agent Initialized[/green]")
    console.print(f"  LLM Tools Prepared: {agent.llm_tools}")


    # 6. Prepare Input and Run Agent
    # The user query isn't strictly needed here as the agent's main task
    # is defined by the system prompt and triggered by the file_path input.
    input_data = ExcelSummarizerInput(file_path=file_path)
    console.print(f"\nRequesting agent to read Excel file: [cyan]{file_path}[/cyan]")
    console.print("[yellow]Agent working... (will attempt to use excel_file_reader tool)[/yellow]")

    try:
        # IMPORTANT: BaseAgent.run() needs modification to handle the tool execution loop.
        # This call will likely just pass the tool info to the LLM.
        # The LLM might respond asking to use the tool, or output a tool call request.
        # This example primarily shows TOOL SETUP. Full execution requires agent logic changes.
        result = agent.run(input_data)

        # Process result (likely won't be ExcelReaderOutput directly from current BaseAgent)
        if isinstance(result, ExcelReaderOutput): # Check if by chance it returned tool output schema
            console.print(Panel(f"[bold green]Tool Output Received (via Agent):[/bold green]", title="Tool Result", expand=False))
            console.print(f"  File: {result.file_path}")
            console.print(f"  Sheet: {result.sheet_name_read}")
            console.print(f"  Rows Read: {result.row_count}")
            console.print(f"  Columns: {result.column_names}")
            console.print(Panel(result.data_preview or "No data preview available.", title="Data Preview"))
        elif isinstance(result, AgentErrorSchema):
             console.print(f"[bold red]Agent Error:[/bold red] {result.error_type} - {result.error_message}")
             if result.details:
                 console.print(f"  Details: {result.details}")
        elif hasattr(result, 'response_message'): # Handle standard text response
             console.print(f"[bold yellow]Agent's Text Response (Tool not executed by agent):[/bold yellow] {result.response_message}")
        else:
            console.print(f"[bold red]Unexpected result type:[/bold red] {type(result)}")
            console.print(result)

    except Exception as e:
        console.print(f"[bold red]An unexpected error occurred:[/bold red] {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Read and preview an Excel file using a Karo agent.")
    parser.add_argument("-f", "--file", required=True, help="Path to the Excel file (.xlsx or .xls).")
    args = parser.parse_args()

    # Check if file exists before running
    if not os.path.exists(args.file):
        console.print(f"[bold red]Error:[/bold red] Input file not found: '{args.file}'")
    else:
        run_summarization(args.file)