""" Financial Document Analysis Workflow - Agno Workflow 2.0 Implementation (Fixed) This workflow processes financial documents through a multi-agent system using the new step-based architecture introduced in Agno Workflow 2.0: 1. Data Extractor Agent: Extracts structured financial data 2. Data Arrangement Function: Organizes data into Excel-ready format 3. Code Generator Agent: Creates professional Excel reports Built according to Agno Workflow 2.0 standards with simple sequential execution. """ import json import time from pathlib import Path from typing import Optional, Dict, Any from textwrap import dedent import os from agno.agent import Agent from agno.models.google import Gemini from agno.tools.file import FileTools from agno.tools.shell import ShellTools from agno.tools.python import PythonTools from agno.workflow.v2.workflow import Workflow from agno.workflow.v2.types import StepInput, StepOutput from agno.workflow.v2.step import Step from agno.storage.sqlite import SqliteStorage # Added this import from agno.utils.log import logger from pydantic import BaseModel, Field from config.settings import settings from utils.prompt_loader import prompt_loader from utils.shell_toolkit import RestrictedShellTools from utils.restricted_python_tools import RestrictedPythonTools class DataPoint(BaseModel): """Individual financial data point.""" field_name: str = Field(description="Name of the financial data field") value: str = Field(description="Value of the field") category: str = Field(description="Financial category (revenue, expenses, assets, etc.)") period: str = Field(default="", description="Time period if applicable") unit: str = Field(default="", description="Currency or measurement unit") confidence: float = Field(default=0.9, description="Confidence score 0-1") class Metadata(BaseModel): """Metadata for extracted financial data.""" company_name: str = Field(default="Unknown Company", description="Company name") document_type: str = Field(default="Unknown", description="Type of financial document") reporting_period: str = Field(default="", description="Reporting period") currency: str = Field(default="", description="Primary currency used") class ExtractedFinancialData(BaseModel): """Structured model for extracted financial data.""" data_points: list[DataPoint] = Field(description="List of extracted financial data points") summary: str = Field(description="Summary of the extracted data") metadata: Metadata = Field(default_factory=Metadata, description="Additional metadata") class FinancialDocumentWorkflow(Workflow): """ Financial document analysis workflow using Agno Workflow 2.0 step-based architecture. This workflow processes financial documents through three specialized steps: - Data extraction with structured outputs - Data arrangement for Excel compatibility - Excel report generation with formatting """ def __init__(self, session_id: Optional[str] = None, **kwargs): """Initialize workflow with session management and step-based architecture.""" # Initialize session directories first self._setup_session_directories(session_id) # Create storage with auto schema upgrade storage = SqliteStorage( table_name="financial_workflows", db_file="tmp/agno_workflows.db", mode="workflow_v2", # Use workflow_v2 mode auto_upgrade_schema=True # This will fix your schema issues ) # Create agents for the workflow self.data_extractor = self._create_data_extractor() self.data_arranger = self._create_data_arranger() self.code_generator = self._create_code_generator() # Create steps using Step objects for better tracking data_extraction_step = Step( name="FinancialDataExtractor", agent=self.data_extractor, description="Expert financial data extraction specialist optimized for Gemini" ) data_arrangement_step = Step( name="DataArrangement", executor=self._arrangement_function, description="User-defined callable step for data arrangement" ) excel_generation_step = Step( name="ExcelReportGenerator", agent=self.code_generator, description="Excel report generator optimized for Gemini with cross-platform support" ) # Initialize the Workflow 2.0 with step-based architecture super().__init__( name="FinancialDocumentWorkflow", description=dedent("""\ Financial document analysis workflow using Agno Workflow 2.0 with step-based execution. Processes financial documents through extraction, arrangement, and Excel report generation. Uses session state for caching and proper error recovery mechanisms. """), steps=[ data_extraction_step, data_arrangement_step, excel_generation_step ], session_id=session_id, storage=storage, # Add the storage here debug_mode=True, **kwargs ) logger.info(f"FinancialDocumentWorkflow v2.0 initialized with session: {self.session_id}") logger.info(f"Session directories: {list(self.session_directories.keys())}") def _setup_session_directories(self, session_id: Optional[str] = None): """Setup session-specific directories.""" self.session_id = session_id self.session_directories = settings.create_session_directories(self.session_id) self.session_output_dir = self.session_directories["output"] self.session_input_dir = self.session_directories["input"] self.session_temp_dir = self.session_directories["temp"] self.session_cache_dir = self.session_directories["cache"] def _create_data_extractor(self) -> Agent: """Create the data extraction agent.""" return Agent( model=Gemini( id=settings.DATA_EXTRACTOR_MODEL, thinking_budget=settings.DATA_EXTRACTOR_MODEL_THINKING_BUDGET, api_key=settings.GOOGLE_API_KEY ), name="FinancialDataExtractor", description="Expert financial data extraction specialist optimized for Gemini", instructions=prompt_loader.load_instructions_as_list("agents/data_extractor"), response_model=ExtractedFinancialData, structured_outputs=True, debug_mode=True, retries=10, delay_between_retries=10, exponential_backoff=True, ) def _create_data_arranger(self) -> Agent: """Create the data arrangement agent.""" logger.info(f"Data arranger base directory: {self.session_output_dir}") logger.info(f"Directory exists: {self.session_output_dir.exists()}") logger.info(f"Directory is writable: {os.access(self.session_output_dir, os.W_OK)}") return Agent( model=Gemini( id=settings.DATA_ARRANGER_MODEL, thinking_budget=settings.DATA_ARRANGER_MODEL_THINKING_BUDGET, api_key=settings.GOOGLE_API_KEY ), name="FinancialDataArranger", description="Financial data organization specialist optimized for Gemini", instructions=prompt_loader.load_instructions_as_list("agents/data_arranger"), tools=[ RestrictedShellTools(base_dir=self.session_output_dir), FileTools(base_dir=self.session_output_dir, save_files=True, read_files=True, list_files=True), ], markdown=False, debug_mode=True, add_memory_references=True, add_session_summary_references=True, retries=10, delay_between_retries=10, exponential_backoff=True, ) def _create_code_generator(self) -> Agent: """Create the code generation agent.""" return Agent( model=Gemini( id=settings.CODE_GENERATOR_MODEL, thinking_budget=settings.CODE_GENERATOR_MODEL_THINKING_BUDGET, api_key=settings.GOOGLE_API_KEY ), name="ExcelReportGenerator", description="Excel report generator optimized for Gemini with cross-platform support", goal="Generate professional Excel reports from arranged financial data with multiple worksheets and formatting", instructions=prompt_loader.load_instructions_as_list("agents/code_generator"), expected_output="A professionally formatted Excel file with multiple worksheets, charts, and proper styling", additional_context=f"Working directory: {self.session_output_dir}. All files must be saved in this directory only.", tools=[ RestrictedShellTools(base_dir=self.session_output_dir), RestrictedPythonTools(base_dir=self.session_output_dir), FileTools(base_dir=self.session_output_dir, save_files=True, read_files=True, list_files=True) ], markdown=False, show_tool_calls=True, debug_mode=True, add_datetime_to_instructions=True, retries=10, delay_between_retries=10, exponential_backoff=True, ) def _arrangement_function(self, step_input: StepInput) -> StepOutput: """Custom function for data arrangement step.""" try: message = step_input.message previous_step_content = step_input.previous_step_content logger.info("Starting data arrangement step") # Load the base arrangement prompt arrangement_prompt = prompt_loader.load_prompt("workflow/data_arrangement") # Combine prompt with extracted data from previous step full_arrangement_prompt = f"{arrangement_prompt}\n\nHere is the extracted financial data to arrange:\n\n{previous_step_content}" # Run data arrangement using the agent response = self.data_arranger.run(full_arrangement_prompt) # Cache the arrangement results in workflow session state if hasattr(self, 'session_state') and self.session_state: cache_key = f"arrangement_{int(time.time())}" self.session_state[cache_key] = response.content logger.info(f"Cached arrangement results with key: {cache_key}") logger.info("Data arrangement completed successfully") return StepOutput( content=response.content, response=response, success=True ) except Exception as e: logger.error(f"Data arrangement failed: {str(e)}") return StepOutput( content=f"Data arrangement failed: {str(e)}", success=False, ) def run(self, file_path: str = None, **kwargs): """ Main workflow execution using Workflow 2.0 step-based architecture. Args: file_path: Path to the financial document to process **kwargs: Additional parameters Returns: Workflow execution result using the new step-based system """ # Handle file_path from kwargs if not provided as positional if file_path is None: file_path = kwargs.get('file_path') if file_path is None: logger.error("file_path is required but not provided") raise ValueError("file_path is required but not provided") start_time = time.time() try: # Validate input file file_path = Path(file_path).resolve() if not file_path.exists(): logger.error(f"File not found: {file_path}") raise FileNotFoundError(f"File not found: {file_path}") # Copy input file to session directory for reference input_file = self.session_input_dir / file_path.name input_file.write_bytes(file_path.read_bytes()) logger.info(f"Starting financial document analysis for: {file_path.name}") # Create File object for direct upload to Gemini API (for first step) from agno.media import File document = File(filepath=str(file_path)) # Load extraction prompt for the first step extraction_prompt = prompt_loader.load_prompt( "workflow/data_extraction", file_path=str(file_path), output_directory=str(self.session_output_dir) ) # Execute the workflow using the new 2.0 step-based system # Pass the extraction prompt as the message and include the file result = super().run( message=extraction_prompt, files=[document], **kwargs ) # Final status execution_time = time.time() - start_time status = self._get_workflow_status() logger.info(f"Workflow completed successfully in {execution_time:.2f} seconds") logger.info(f"Results: {status}") return result except Exception as e: logger.error(f"Workflow execution failed: {str(e)}") raise def _get_workflow_status(self) -> Dict[str, Any]: """Get current workflow status and file counts.""" status = { "session_id": self.session_id, "output_directory": str(self.session_output_dir), "json_files": 0, "excel_files": 0, "data_points": 0 } if self.session_output_dir.exists(): status["json_files"] = len(list(self.session_output_dir.glob("*.json"))) status["excel_files"] = len(list(self.session_output_dir.glob("*.xlsx"))) return status # Compatibility function to maintain the same interface as the original workflow def create_financial_workflow(session_id: Optional[str] = None, **kwargs) -> FinancialDocumentWorkflow: """ Create a new FinancialDocumentWorkflow instance using Workflow 2.0. Args: session_id: Optional session ID for tracking workflow execution **kwargs: Additional parameters for workflow configuration Returns: FinancialDocumentWorkflow: Configured workflow instance """ return FinancialDocumentWorkflow(session_id=session_id, **kwargs)