Session 1: From Prototype to Production

"""
Multi-Agent Workflow - Complete Solution
Production-ready implementation with error handling,
cost tracking, structured logging, and monitoring.
"""

import os
import json
import time
import logging
from typing import Dict, List, TypedDict, Optional
from dataclasses import dataclass, field
from datetime import datetime
import traceback

import anthropic

# Configure structured logging
logging.basicConfig(
    level=logging.INFO,
    format='%(message)s'
)
logger = logging.getLogger(__name__)

def log_structured(event: str, **kwargs):
    """Structured JSON logging for production observability"""
    log_entry = {
        "timestamp": datetime.utcnow().isoformat(),
        "event": event,
        **kwargs
    }
    logger.info(json.dumps(log_entry))

# Configuration
class Config:
    ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
    MODEL = "claude-3-5-sonnet-20241022"
    MAX_TOKENS = 2048
    TIMEOUT = 30  # seconds
    MAX_RETRIES = 3
    RETRY_DELAY = 2  # seconds

    # Cost tracking (per 1M tokens)
    COST_PER_INPUT_TOKEN = 0.003 / 1000
    COST_PER_OUTPUT_TOKEN = 0.015 / 1000

# Workflow State
class WorkflowState(TypedDict):
    topic: str
    research_data: str
    analysis: str
    final_report: str
    error: Optional[str]
    metadata: Dict

# Cost Tracker
@dataclass
class CostTracker:
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    total_cost: float = 0.0
    agent_costs: Dict[str, float] = field(default_factory=dict)

    def add_usage(self, agent_name: str, input_tokens: int, output_tokens: int):
        cost = (
            input_tokens * Config.COST_PER_INPUT_TOKEN +
            output_tokens * Config.COST_PER_OUTPUT_TOKEN
        )
        self.total_input_tokens += input_tokens
        self.total_output_tokens += output_tokens
        self.total_cost += cost

        if agent_name not in self.agent_costs:
            self.agent_costs[agent_name] = 0.0
        self.agent_costs[agent_name] += cost

        log_structured(
            "cost_tracking",
            agent=agent_name,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cost_usd=round(cost, 4)
        )

    def get_report(self) -> Dict:
        return {
            "total_input_tokens": self.total_input_tokens,
            "total_output_tokens": self.total_output_tokens,
            "total_cost_usd": round(self.total_cost, 4),
            "agent_breakdown": {
                agent: round(cost, 4)
                for agent, cost in self.agent_costs.items()
            }
        }

# Global cost tracker
cost_tracker = CostTracker()

# Base Agent with production features
class BaseAgent:
    def __init__(self, name: str, config: Config):
        self.name = name
        self.config = config
        self.client = anthropic.Anthropic(
            api_key=config.ANTHROPIC_API_KEY
        )

    def call_llm(
        self,
        system_prompt: str,
        user_message: str,
        max_retries: int = None
    ) -> str:
        """
        Production LLM call with:
        - Retry logic with exponential backoff
        - Timeout handling
        - Error recovery
        - Cost tracking
        """
        retries = max_retries or self.config.MAX_RETRIES

        for attempt in range(retries):
            try:
                log_structured(
                    "llm_call_start",
                    agent=self.name,
                    attempt=attempt + 1,
                    max_retries=retries
                )

                start_time = time.time()

                response = self.client.messages.create(
                    model=self.config.MODEL,
                    max_tokens=self.config.MAX_TOKENS,
                    system=system_prompt,
                    messages=[{
                        "role": "user",
                        "content": user_message
                    }],
                    timeout=self.config.TIMEOUT
                )

                duration = time.time() - start_time

                # Track costs
                cost_tracker.add_usage(
                    self.name,
                    response.usage.input_tokens,
                    response.usage.output_tokens
                )

                # Extract response
                result = response.content[0].text

                log_structured(
                    "llm_call_success",
                    agent=self.name,
                    duration_seconds=round(duration, 2),
                    output_length=len(result)
                )

                return result

            except anthropic.APITimeoutError as e:
                log_structured(
                    "llm_call_timeout",
                    agent=self.name,
                    attempt=attempt + 1,
                    error=str(e)
                )

                if attempt < retries - 1:
                    wait_time = self.config.RETRY_DELAY * (2 ** attempt)
                    time.sleep(wait_time)
                    continue
                else:
                    raise

            except anthropic.RateLimitError as e:
                log_structured(
                    "llm_call_rate_limit",
                    agent=self.name,
                    attempt=attempt + 1
                )

                if attempt < retries - 1:
                    wait_time = self.config.RETRY_DELAY * (2 ** attempt)
                    time.sleep(wait_time)
                    continue
                else:
                    raise

            except Exception as e:
                log_structured(
                    "llm_call_error",
                    agent=self.name,
                    attempt=attempt + 1,
                    error=str(e),
                    traceback=traceback.format_exc()
                )

                if attempt < retries - 1:
                    time.sleep(self.config.RETRY_DELAY)
                    continue
                else:
                    raise

        raise Exception(f"Max retries ({retries}) exceeded")

# Research Agent
class ResearchAgent(BaseAgent):
    def __init__(self, config: Config):
        super().__init__("researcher", config)

    def research(self, topic: str) -> str:
        """Gather research data on topic"""
        log_structured("agent_start", agent=self.name, topic=topic)

        system_prompt = """You are a professional research agent.
Your task is to gather comprehensive information on the given topic.

Focus on:
1. Key facts and data points
2. Current trends and developments
3. Expert opinions and analysis
4. Relevant statistics

Provide well-structured research findings in 3-4 paragraphs."""

        user_message = f"Research the following topic:\n\n{topic}"

        try:
            result = self.call_llm(system_prompt, user_message)
            log_structured("agent_complete", agent=self.name, success=True)
            return result
        except Exception as e:
            log_structured("agent_failed", agent=self.name, error=str(e))
            return f"Research failed: {str(e)}"

# Analyst Agent
class AnalystAgent(BaseAgent):
    def __init__(self, config: Config):
        super().__init__("analyst", config)

    def analyze(self, research_data: str) -> str:
        """Analyze research data"""
        log_structured("agent_start", agent=self.name)

        system_prompt = """You are an expert analyst.
Your task is to analyze the provided research data and extract insights.

Focus on:
1. Key patterns and trends
2. Implications and consequences
3. Opportunities and risks
4. Strategic recommendations

Provide structured analysis in 3-4 paragraphs."""

        user_message = f"Analyze this research data:\n\n{research_data}"

        try:
            result = self.call_llm(system_prompt, user_message)
            log_structured("agent_complete", agent=self.name, success=True)
            return result
        except Exception as e:
            log_structured("agent_failed", agent=self.name, error=str(e))
            return f"Analysis failed: {str(e)}"

# Writer Agent
class WriterAgent(BaseAgent):
    def __init__(self, config: Config):
        super().__init__("writer", config)

    def write_report(self, analysis: str, topic: str) -> str:
        """Write final report"""
        log_structured("agent_start", agent=self.name)

        system_prompt = """You are a professional report writer.
Your task is to create a clear, compelling executive report.

Format:
1. Executive Summary (2-3 sentences)
2. Key Findings (bullet points)
3. Detailed Analysis (2-3 paragraphs)
4. Recommendations (bullet points)

Use clear, professional language suitable for decision-makers."""

        user_message = f"""Create an executive report on: {topic}

Based on this analysis:
{analysis}"""

        try:
            result = self.call_llm(system_prompt, user_message)
            log_structured("agent_complete", agent=self.name, success=True)
            return result
        except Exception as e:
            log_structured("agent_failed", agent=self.name, error=str(e))
            return f"Report writing failed: {str(e)}"

# Workflow Orchestration
class MultiAgentWorkflow:
    def __init__(self, config: Config):
        self.config = config
        self.researcher = ResearchAgent(config)
        self.analyst = AnalystAgent(config)
        self.writer = WriterAgent(config)

    def execute(self, topic: str) -> Dict:
        """Execute full workflow with error handling"""
        log_structured("workflow_start", topic=topic)
        start_time = time.time()

        state = {
            "topic": topic,
            "research_data": "",
            "analysis": "",
            "final_report": "",
            "error": None,
            "metadata": {
                "start_time": datetime.utcnow().isoformat(),
                "status": "running"
            }
        }

        try:
            # Step 1: Research
            log_structured("workflow_step", step="research")
            state["research_data"] = self.researcher.research(topic)

            if "failed" in state["research_data"].lower():
                raise Exception("Research step failed")

            # Step 2: Analysis
            log_structured("workflow_step", step="analysis")
            state["analysis"] = self.analyst.analyze(state["research_data"])

            if "failed" in state["analysis"].lower():
                raise Exception("Analysis step failed")

            # Step 3: Writing
            log_structured("workflow_step", step="writing")
            state["final_report"] = self.writer.write_report(
                state["analysis"],
                topic
            )

            if "failed" in state["final_report"].lower():
                raise Exception("Writing step failed")

            # Success
            duration = time.time() - start_time
            state["metadata"].update({
                "status": "completed",
                "duration_seconds": round(duration, 2),
                "end_time": datetime.utcnow().isoformat(),
                "costs": cost_tracker.get_report()
            })

            log_structured(
                "workflow_complete",
                success=True,
                duration_seconds=round(duration, 2)
            )

        except Exception as e:
            duration = time.time() - start_time
            state["error"] = str(e)
            state["metadata"].update({
                "status": "failed",
                "duration_seconds": round(duration, 2),
                "end_time": datetime.utcnow().isoformat(),
                "error": str(e)
            })

            log_structured(
                "workflow_failed",
                error=str(e),
                traceback=traceback.format_exc()
            )

        return state

# Main Execution
if __name__ == "__main__":
    print("=== Multi-Agent Workflow Demo ===\n")

    # Create workflow
    workflow = MultiAgentWorkflow(Config)

    # Test topics
    topics = [
        "Impact of AI on software development productivity",
        "Future of remote work in tech companies",
        "Quantum computing applications in finance"
    ]

    print("Available topics:")
    for i, topic in enumerate(topics, 1):
        print(f"  {i}. {topic}")

    choice = input("\nChoose topic (1-3) or enter custom: ").strip()

    if choice.isdigit() and 1 <= int(choice) <= len(topics):
        topic = topics[int(choice) - 1]
    else:
        topic = choice if choice else topics[0]

    print(f"\n{'='*60}")
    print(f"Processing: {topic}")
    print(f"{'='*60}\n")

    # Execute workflow
    result = workflow.execute(topic)

    # Display results
    if result["error"]:
        print(f"\n❌ Workflow failed: {result['error']}\n")
    else:
        print("\n" + "="*60)
        print("FINAL REPORT")
        print("="*60 + "\n")
        print(result["final_report"])
        print("\n" + "="*60)

    # Display metrics
    print("\n" + "="*60)
    print("WORKFLOW METRICS")
    print("="*60)
    print(f"Status: {result['metadata']['status']}")
    print(f"Duration: {result['metadata']['duration_seconds']}s")

    if "costs" in result["metadata"]:
        costs = result["metadata"]["costs"]
        print(f"\nCost Breakdown:")
        print(f"  Total: ${costs['total_cost_usd']}")
        print(f"  Input tokens: {costs['total_input_tokens']}")
        print(f"  Output tokens: {costs['total_output_tokens']}")
        print(f"\n  By Agent:")
        for agent, cost in costs['agent_breakdown'].items():
            print(f"    {agent}: ${cost}")

    print("\n=== Demo Complete ===")