agents/simple_agent.py

"""
Simplified financial agent focused only on generating financial statements
"""
from typing import Dict, Any, Optional
import logging
from .base_config import FinRyverAgentConfig, create_agent_executor
from .simple_tools import FINANCIAL_TOOLS

logger = logging.getLogger(__name__)

class FinancialStatementAgent:
    """
    Simple agent that generates: Notes, Balance Sheet, P&L, Cash Flow
    """
    
    def __init__(self, config: Optional[FinRyverAgentConfig] = None):
        self.config = config or FinRyverAgentConfig()
        self.llm = self.config.get_llm()
        
        # Create agent with financial statement tools
        self.agent_executor = create_agent_executor(
            role="Financial Statement Generator",
            goal="Generate financial statements (notes, balance sheet, P&L, cash flow) from trial balance data",
            tools=FINANCIAL_TOOLS,
            context="You generate financial statements from trial balance files. Available outputs: Notes, Balance Sheet, P&L Statement, Cash Flow Statement.",
            llm=self.llm
        )
    
    
    async def generate_specific_statement(self, file_path: str, statement_type: str, note_numbers: str) -> Dict[str, Any]:
        """Generate a specific financial statement"""
        try:
            prompts = {
                "notes": f"Generate financial notes for note numbers {note_numbers} from {file_path}",
                "notes-llm": f"Generate financial notes using LLM for note numbers {note_numbers} from {file_path}",
                "balance_sheet": f"Generate balance sheet from {file_path} using note numbers {note_numbers}",
                "pnl": f"Generate P&L statement from {file_path} using note numbers {note_numbers}",
                "cash_flow": f"Generate cash flow statement from {file_path} using note numbers {note_numbers}"
            }
            
            prompt = prompts.get(statement_type, prompts["notes"])
            
            result = await self.agent_executor.ainvoke({"input": prompt})
            return {"status": "success", "result": result["output"]}
            
        except Exception as e:
            logger.error(f"Error generating {statement_type}: {e}")
            return {"status": "error", "error": str(e)}