examples/integration_example.py

"""Example integration of Agent Cost Optimizer with a hypothetical agent harness."""

from aco import AgentCostOptimizer
from aco.config import ACOConfig


def example_agent_harness():
    """Example of how to bolt ACO onto any agent harness."""
    
    # Initialize optimizer
    config = ACOConfig.from_yaml("config.yaml")
    optimizer = AgentCostOptimizer(config)
    
    # Incoming user request
    user_request = "Write a Python script to fetch data from an API and cache it in Redis"
    
    # Build run state from current agent state
    run_state = {
        "trace_id": "agent-run-12345",
        "current_cost": 0.0,
        "planned_tools": [
            ("search", {"query": "redis python client"}),
            ("fetch", {"url": "https://api.example.com/docs"}),
            ("code_execution", {"code": "test script"}),
        ],
        "previous_tool_calls": [],
        "step_number": 1,
        "total_steps": 3,
        "is_irreversible": False,
        "context_pieces": {
            "system_rules": "You are a coding assistant.",
            "tool_descriptions": "Available tools: search, fetch, code_execution",
            "user_preferences": "Prefer Python 3.11+, type hints, async where possible",
            "recent_messages": "User: Write a Python script...",
        },
        "retrieved_docs": [],
        "routing_mode": "cascade",
    }
    
    # Call optimizer before executing
    decision = optimizer.optimize(user_request, run_state)
    
    print(f"Trace ID: {decision.trace_id}")
    print(f"Selected Model: {decision.routing_decision.model_id} (tier {decision.routing_decision.tier})")
    print(f"Estimated Cost: ${decision.estimated_cost:.4f}")
    print(f"Estimated Latency: {decision.estimated_latency_ms:.0f}ms")
    print(f"Confidence: {decision.confidence:.2f}")
    print()
    
    # Apply tool gate decisions
    print("Tool Decisions:")
    for td in decision.tool_decisions:
        print(f"  {td.tool_name}: {td.decision.value} (reason: {td.reasoning})")
    
    # Apply context budget
    if decision.context_budget:
        print(f"\nContext Budget: {decision.context_budget.total_budget_tokens} tokens")
        print(f"  Cache prefix: {decision.context_budget.cache_prefix_tokens} tokens")
        print(f"  Dynamic suffix: {decision.context_budget.dynamic_suffix_tokens} tokens")
        if decision.context_budget.omitted_sources:
            print(f"  Omitted: {[s.name for s in decision.context_budget.omitted_sources]}")
    
    # Apply cache layout
    if decision.prompt_layout:
        print(f"\nCache Layout:")
        print(f"  Cold cost: ${decision.prompt_layout.estimated_cold_cost:.4f}")
        print(f"  Warm cost: ${decision.prompt_layout.estimated_warm_cost:.4f}")
        print(f"  Cache discount: ${decision.prompt_layout.cache_discount:.4f}")
    
    # Check meta-tool
    if decision.meta_tool_match:
        print(f"\nMeta-Tool Match: {decision.meta_tool_match['meta_tool_id']}")
        print(f"  Estimated savings: ${decision.meta_tool_match['estimated_cost_savings']:.4f}")
    
    # Check doom assessment
    if decision.doom_assessment:
        print(f"\nDoom Assessment: {decision.doom_assessment.action.value}")
        print(f"  Confidence: {decision.doom_assessment.confidence:.2f}")
        print(f"  Signals: {decision.doom_assessment.signals_triggered}")
    
    # Check verifier
    if decision.verifier_decision:
        print(f"\nVerifier: {decision.verifier_decision.decision.value}")
        print(f"  Checks: {decision.verifier_decision.checks}")
        print(f"  Cost: ${decision.verifier_decision.estimated_verifier_cost:.4f}")
    
    # After execution, record step and finalize
    from aco.trace_schema import ModelCall, Outcome
    
    model_call = ModelCall(
        model_id=decision.routing_decision.model_id,
        provider="cloud",
        input_tokens=2048,
        output_tokens=512,
        cost_per_1k_input=0.003,
        cost_per_1k_output=0.006,
    )
    
    optimizer.record_step(
        trace_id=decision.trace_id,
        model_call=model_call,
        context_size_tokens=2048,
        step_outcome=Outcome.SUCCESS,
    )
    
    # Finalize
    trace = optimizer.finalize_trace(
        trace_id=decision.trace_id,
        outcome=Outcome.SUCCESS,
        user_satisfaction=0.95,
    )
    
    print(f"\nTrace finalized. Total cost: ${trace.total_cost_computed:.4f}")
    print(f"Cost saved vs frontier: ${trace.total_cost_saved_vs_frontier:.4f}")


if __name__ == "__main__":
    example_agent_harness()