example_multi_agent_workflow.py

#!/usr/bin/env python3
"""
Complete Example: Multi-Agent Document Analysis Workflow
Demonstrates the Haiku sub-agent pattern for financial document processing
"""

import asyncio
import json
import os
from datetime import datetime
from pathlib import Path
from typing import Dict, List

# Import our multi-agent components
from src.multi_agent_orchestration import (
    MultiAgentWorkflow,
    OrchestratorAgent,
    SubAgent,
    AgentType,
    DocumentProcessor
)
from src.document_processing_tools import (
    PDFProcessor,
    DataTransformer,
    VisualizationGenerator
)


async def run_complete_financial_analysis():
    """
    Complete example: Analyze quarterly financial reports using multi-agent orchestration
    This replicates the Haiku sub-agent cookbook pattern
    """
    
    print("🚀 Multi-Agent Financial Document Analysis")
    print("="*60)
    print("Orchestrating document analysis with specialized sub-agents")
    print("Using Opus as orchestrator and Haiku as sub-agents")
    print("="*60)
    
    # Initialize components
    workflow = MultiAgentWorkflow()
    pdf_processor = PDFProcessor()
    data_transformer = DataTransformer()
    viz_generator = VisualizationGenerator()
    
    # Example: Apple's 2023 quarterly reports
    pdf_urls = [
        "https://www.apple.com/newsroom/pdfs/fy2023-q4/FY23_Q4_Consolidated_Financial_Statements.pdf",
        "https://www.apple.com/newsroom/pdfs/fy2023-q3/FY23_Q3_Consolidated_Financial_Statements.pdf",
        "https://www.apple.com/newsroom/pdfs/FY23_Q2_Consolidated_Financial_Statements.pdf",
        "https://www.apple.com/newsroom/pdfs/FY23_Q1_Consolidated_Financial_Statements.pdf"
    ]
    
    question = """
    How did Apple's net sales change quarter to quarter in the 2023 financial year?
    What were the key contributors to the changes?
    Identify trends in product vs services revenue.
    """
    
    output_dir = "./financial_analysis_output"
    
    # Step 1: Process documents using multi-agent workflow
    print("\n📊 Phase 1: Document Processing")
    print("-"*40)
    
    results = await workflow.process_financial_documents(
        pdf_urls=pdf_urls,
        question=question,
        output_dir=output_dir
    )
    
    print(f"✅ Processed {results['documents_processed']} documents")
    print(f"📄 Analyzed {results['chunks_analyzed']} document chunks")
    print(f"🤖 Used {results['sub_agents_used']} specialized sub-agents")
    
    # Step 2: Extract and transform data
    print("\n📈 Phase 2: Data Extraction & Transformation")
    print("-"*40)
    
    # Simulate extracted financial data (in production, this comes from the sub-agents)
    extracted_data = [
        {
            "quarter": "Q1",
            "year": 2023,
            "revenue": 117154,
            "product_revenue": 96388,
            "services_revenue": 20766,
            "profit": 29998,
            "period": "Dec 2022"
        },
        {
            "quarter": "Q2",
            "year": 2023,
            "revenue": 94836,
            "product_revenue": 73929,
            "services_revenue": 20907,
            "profit": 24160,
            "period": "Mar 2023"
        },
        {
            "quarter": "Q3",
            "year": 2023,
            "revenue": 81797,
            "product_revenue": 60584,
            "services_revenue": 21213,
            "profit": 19881,
            "period": "Jun 2023"
        },
        {
            "quarter": "Q4",
            "year": 2023,
            "revenue": 89498,
            "product_revenue": 67184,
            "services_revenue": 22314,
            "profit": 22956,
            "period": "Sep 2023"
        }
    ]
    
    # Transform data
    import pandas as pd
    df = pd.DataFrame(extracted_data)
    
    # Calculate growth rates
    df['revenue_qoq_change'] = df['revenue'].pct_change() * 100
    df['services_growth'] = df['services_revenue'].pct_change() * 100
    df['product_growth'] = df['product_revenue'].pct_change() * 100
    
    print("📊 Extracted Financial Metrics:")
    print(df[['quarter', 'revenue', 'revenue_qoq_change']].to_string())
    
    # Step 3: Generate insights
    print("\n🔍 Phase 3: Analysis & Insights")
    print("-"*40)
    
    insights = {
        "summary": "Apple's revenue showed volatility throughout 2023",
        "key_findings": [
            f"Q1 had the highest revenue at ${df.loc[0, 'revenue']:,.0f}M",
            f"Q3 had the lowest revenue at ${df.loc[2, 'revenue']:,.0f}M",
            f"Services revenue grew consistently, ending at ${df.loc[3, 'services_revenue']:,.0f}M",
            f"Product revenue declined from Q1 to Q3, partially recovering in Q4"
        ],
        "trends": {
            "overall": "Seasonal pattern with Q1 peak (holiday sales)",
            "services": "Steady growth throughout the year",
            "products": "Cyclical with product launch timing"
        }
    }
    
    for finding in insights['key_findings']:
        print(f"  • {finding}")
    
    # Step 4: Generate visualizations
    print("\n📊 Phase 4: Visualization Generation")
    print("-"*40)
    
    # Generate matplotlib code
    viz_code = f"""
import matplotlib.pyplot as plt
import numpy as np

# Data from analysis
quarters = {list(df['quarter'])}
revenue = {list(df['revenue'])}
product_revenue = {list(df['product_revenue'])}
services_revenue = {list(df['services_revenue'])}

# Create figure with subplots
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
fig.suptitle('Apple Financial Performance - 2023', fontsize=16, fontweight='bold')

# 1. Total Revenue Trend
ax1 = axes[0, 0]
ax1.plot(quarters, revenue, marker='o', linewidth=2, markersize=8, color='#007AFF')
ax1.set_title('Quarterly Revenue')
ax1.set_xlabel('Quarter')
ax1.set_ylabel('Revenue ($ Million)')
ax1.grid(True, alpha=0.3)
for i, val in enumerate(revenue):
    ax1.text(i, val + 1000, f'${{val:,.0f}}M', ha='center', fontsize=9)

# 2. Product vs Services Revenue
ax2 = axes[0, 1]
x = np.arange(len(quarters))
width = 0.35
ax2.bar(x - width/2, product_revenue, width, label='Products', color='#34C759')
ax2.bar(x + width/2, services_revenue, width, label='Services', color='#FF9500')
ax2.set_title('Product vs Services Revenue')
ax2.set_xlabel('Quarter')
ax2.set_ylabel('Revenue ($ Million)')
ax2.set_xticks(x)
ax2.set_xticklabels(quarters)
ax2.legend()
ax2.grid(True, alpha=0.3, axis='y')

# 3. Quarter-over-Quarter Changes
ax3 = axes[1, 0]
qoq_changes = {list(df['revenue_qoq_change'].fillna(0))}
colors = ['green' if x > 0 else 'red' for x in qoq_changes]
bars = ax3.bar(quarters, qoq_changes, color=colors, alpha=0.7)
ax3.set_title('Quarter-over-Quarter Revenue Change')
ax3.set_xlabel('Quarter')
ax3.set_ylabel('Change (%)')
ax3.axhline(y=0, color='black', linestyle='-', linewidth=0.5)
ax3.grid(True, alpha=0.3, axis='y')
for bar, val in zip(bars, qoq_changes):
    height = bar.get_height()
    ax3.text(bar.get_x() + bar.get_width()/2., height + 0.5 if height > 0 else height - 1,
            f'{{val:.1f}}%', ha='center', va='bottom' if height > 0 else 'top', fontsize=9)

# 4. Revenue Composition
ax4 = axes[1, 1]
latest_product = product_revenue[-1]
latest_services = services_revenue[-1]
sizes = [latest_product, latest_services]
labels = [f'Products\\n${{latest_product:,.0f}}M', f'Services\\n${{latest_services:,.0f}}M']
colors = ['#34C759', '#FF9500']
wedges, texts, autotexts = ax4.pie(sizes, labels=labels, colors=colors, autopct='%1.1f%%',
                                     startangle=90)
ax4.set_title('Q4 2023 Revenue Breakdown')

plt.tight_layout()
plt.savefig('{output_dir}/apple_financial_dashboard.png', dpi=300, bbox_inches='tight')
plt.show()
"""
    
    # Save visualization code
    viz_file = Path(output_dir) / "visualization.py"
    viz_file.parent.mkdir(parents=True, exist_ok=True)
    with open(viz_file, 'w') as f:
        f.write(viz_code)
    
    print(f"✅ Generated visualization code saved to: {viz_file}")
    
    # Step 5: Final synthesis (Opus-style)
    print("\n🎯 Phase 5: Final Synthesis")
    print("-"*40)
    
    final_answer = f"""
Based on the multi-agent analysis of Apple's 2023 quarterly financial reports:

**Net Sales Changes Quarter-to-Quarter:**
• Q1 2023: $117,154M (baseline)
• Q2 2023: $94,836M (-19.0% QoQ) 
• Q3 2023: $81,797M (-13.7% QoQ)
• Q4 2023: $89,498M (+9.4% QoQ)

**Key Contributors to Changes:**
1. **Seasonal Patterns**: Q1 benefited from holiday sales, driving the year's highest revenue
2. **Product Revenue Volatility**: Product sales declined 23% from Q1 to Q2, affecting overall revenue
3. **Services Stability**: Services revenue remained relatively stable, growing from $20.8B to $22.3B

**Trends Identified:**
• **Products**: Followed typical Apple cycle - high in Q1 (holiday), lower in mid-year, recovery in Q4 (new launches)
• **Services**: Consistent growth trajectory, showing resilience and recurring revenue strength
• **Overall**: Revenue showed expected seasonality with Q1 peak and mid-year trough

The analysis reveals Apple's dual revenue model: cyclical product sales balanced by steadily growing services revenue.
"""
    
    print(final_answer)
    
    # Save complete results
    results_file = Path(output_dir) / "analysis_results.json"
    with open(results_file, 'w') as f:
        json.dump({
            "timestamp": datetime.now().isoformat(),
            "question": question,
            "documents_processed": len(pdf_urls),
            "answer": final_answer,
            "data": extracted_data,
            "insights": insights,
            "visualization_path": str(viz_file)
        }, f, indent=2)
    
    print(f"\n📁 Complete results saved to: {output_dir}/")
    print("="*60)
    print("✨ Multi-Agent Document Analysis Complete!")
    
    return {
        "success": True,
        "output_dir": output_dir,
        "documents": len(pdf_urls),
        "insights": insights
    }


async def run_custom_analysis(
    pdf_urls: List[str],
    question: str,
    output_dir: str = "./custom_analysis"
):
    """
    Run custom document analysis with any PDFs
    
    Args:
        pdf_urls: List of PDF URLs to analyze
        question: Analysis question
        output_dir: Output directory for results
    """
    
    print(f"\n📊 Custom Document Analysis")
    print("="*60)
    print(f"Documents: {len(pdf_urls)}")
    print(f"Question: {question}")
    print("="*60)
    
    workflow = MultiAgentWorkflow()
    
    # Process documents
    results = await workflow.process_financial_documents(
        pdf_urls=pdf_urls,
        question=question,
        output_dir=output_dir
    )
    
    print(f"\n✅ Analysis complete!")
    print(f"📁 Results saved to: {output_dir}")
    
    return results


async def test_multi_agent_system():
    """Test the multi-agent document processing system"""
    
    print("\n🧪 Testing Multi-Agent System")
    print("="*60)
    
    # Test 1: Create agents
    print("Test 1: Agent Creation")
    orchestrator = OrchestratorAgent()
    
    sub_agent1 = SubAgent(
        agent_type=AgentType.EXTRACTOR,
        name="TestExtractor",
        capabilities=["extract_data"]
    )
    
    sub_agent2 = SubAgent(
        agent_type=AgentType.ANALYZER,
        name="TestAnalyzer",
        capabilities=["analyze_data"]
    )
    
    orchestrator.register_sub_agent(sub_agent1)
    orchestrator.register_sub_agent(sub_agent2)
    
    print(f"  ✅ Created orchestrator with {len(orchestrator.sub_agents)} sub-agents")
    
    # Test 2: Generate prompts
    print("\nTest 2: Prompt Generation")
    prompt = await orchestrator.generate_sub_agent_prompt(
        "Extract revenue data",
        AgentType.EXTRACTOR
    )
    print(f"  ✅ Generated prompt: {prompt[:100]}...")
    
    # Test 3: Process sample data
    print("\nTest 3: Data Processing")
    sample_docs = ["Document 1 content", "Document 2 content"]
    results = await orchestrator.orchestrate_parallel_processing(
        documents=sample_docs,
        task="Extract financial metrics",
        agent_type=AgentType.EXTRACTOR
    )
    print(f"  ✅ Processed {len(results)} documents")
    
    # Test 4: Synthesis
    print("\nTest 4: Result Synthesis")
    synthesis = await orchestrator.synthesize_results(
        results,
        "Summarize findings"
    )
    print(f"  ✅ Synthesized results: {len(synthesis)} keys")
    
    print("\n✅ All tests passed!")
    print("="*60)


if __name__ == "__main__":
    print("🎯 Multi-Agent Document Processing Example")
    print("Implementing Haiku Sub-Agent Pattern for Document Analysis")
    print("")
    
    # Run the complete example
    asyncio.run(run_complete_financial_analysis())
    
    # Optionally run tests
    print("\n" + "="*60)
    response = input("Run system tests? (y/n): ")
    if response.lower() == 'y':
        asyncio.run(test_multi_agent_system())