headroom/examples/multi_tool_agent_test.py

#!/usr/bin/env python3
"""
Multi-Tool Agent Demo: Headroom in Action

Shows an AI agent investigating a memory leak using 4 tools.
Headroom compresses the tool outputs while preserving critical information.

Usage:
    export ANTHROPIC_API_KEY=sk-ant-...
    python examples/multi_tool_agent_test.py
"""

import json
import os
import time

from agno.agent import Agent
from agno.models.anthropic import Claude
from agno.tools import tool

if not os.environ.get("ANTHROPIC_API_KEY"):
    raise ValueError("ANTHROPIC_API_KEY environment variable required")


# =============================================================================
# MOCK TOOL DATA
# Each tool returns many items, but only ONE is critical (the "needle")
# =============================================================================


def generate_github_issues(count: int = 50, needle_at: int = 42) -> list[dict]:
    """Generate GitHub issues with a memory leak bug at needle_at."""
    issues = []
    for i in range(count):
        if i == needle_at:
            issues.append(
                {
                    "number": i,
                    "title": "Memory leak in worker pool",
                    "state": "open",
                    "labels": ["bug", "priority:high"],
                    "body": "Worker threads not released after task completion. Memory grows unboundedly.",
                    "reactions": {"thumbs_up": 47},
                }
            )
        else:
            issues.append(
                {
                    "number": i,
                    "title": f"Feature: {['dark mode', 'pagination', 'webhooks', 'rate limiting'][i % 4]}",
                    "state": "closed" if i % 3 == 0 else "open",
                    "labels": ["enhancement"],
                    "body": "Please add this feature.",
                    "reactions": {"thumbs_up": i % 5},
                }
            )
    return issues


def generate_code_results(count: int = 40, needle_at: int = 23) -> list[dict]:
    """Generate code search results with memory fix at needle_at."""
    results = []
    for i in range(count):
        if i == needle_at:
            results.append(
                {
                    "file": "src/worker.py",
                    "line": 330,
                    "content": """def cleanup_worker(self):
    \"\"\"Release worker resources - MEMORY LEAK FIX\"\"\"
    self.thread_pool.shutdown(wait=True)
    self.connections.clear()
    gc.collect()""",
                    "relevance_score": 0.98,
                }
            )
        else:
            results.append(
                {
                    "file": f"src/{['utils', 'api', 'models', 'handlers'][i % 4]}.py",
                    "line": 100 + i * 10,
                    "content": f"def process_{['data', 'request', 'model', 'event'][i % 4]}(self): pass",
                    "relevance_score": 0.5 - (i * 0.01),
                }
            )
    return results


def generate_db_logs(count: int = 60, needle_at: int = 17) -> list[dict]:
    """Generate database logs with OutOfMemoryError at needle_at."""
    logs = []
    for i in range(count):
        if i == needle_at:
            logs.append(
                {
                    "type": "error",
                    "service": "worker-pool",
                    "message": "OutOfMemoryError: heap space exhausted",
                    "metadata": {"heap_used": "7.8GB", "heap_max": "8GB", "thread_count": 847},
                    "stack_trace": "java.lang.OutOfMemoryError at WorkerPool.execute()",
                }
            )
        else:
            logs.append(
                {
                    "type": "info",
                    "service": ["api", "auth", "db", "cache"][i % 4],
                    "message": "Operation completed successfully",
                    "metadata": {"heap_used": f"{1 + i % 3}GB", "thread_count": 50 + i % 50},
                }
            )
    return logs


def generate_arxiv_papers(count: int = 30, needle_at: int = 15) -> list[dict]:
    """Generate ArXiv papers with relevant research at needle_at."""
    papers = []
    for i in range(count):
        if i == needle_at:
            papers.append(
                {
                    "id": f"2401.{i:05d}",
                    "title": "Memory Management in Worker Pool Architectures",
                    "abstract": "We analyze memory leak patterns in thread pools and propose cleanup strategies.",
                    "citations": 89,
                }
            )
        else:
            papers.append(
                {
                    "id": f"2401.{i:05d}",
                    "title": f"{['Deep Learning', 'Transformers', 'Neural Arch', 'Scaling'][i % 4]} Study",
                    "abstract": "Novel approach with state-of-the-art results.",
                    "citations": i * 3,
                }
            )
    return papers


# Generate the mock data
GITHUB_ISSUES = generate_github_issues()
CODE_RESULTS = generate_code_results()
DB_LOGS = generate_db_logs()
ARXIV_PAPERS = generate_arxiv_papers()


# =============================================================================
# TOOLS
# =============================================================================


@tool(name="search_github_issues")
def search_github_issues(query: str) -> str:
    """Search GitHub issues."""
    return json.dumps(GITHUB_ISSUES, indent=2)


@tool(name="search_code")
def search_code(query: str) -> str:
    """Search codebase."""
    return json.dumps(CODE_RESULTS, indent=2)


@tool(name="query_database")
def query_database(query: str) -> str:
    """Query database logs."""
    return json.dumps(DB_LOGS, indent=2)


@tool(name="search_arxiv")
def search_arxiv(query: str) -> str:
    """Search ArXiv papers."""
    return json.dumps(ARXIV_PAPERS, indent=2)


# =============================================================================
# VERIFICATION
# =============================================================================

NEEDLES = {
    "GitHub Issue #42": lambda r: "memory leak" in r.lower()
    and ("42" in r or "worker" in r.lower()),
    "cleanup_worker() fix": lambda r: "cleanup" in r.lower() or "worker.py" in r.lower(),
    "OutOfMemoryError": lambda r: "outofmemory" in r.lower() or "847" in r or "7.8" in r.lower(),
    "ArXiv paper": lambda r: "paper" in r.lower()
    or "research" in r.lower()
    or "arxiv" in r.lower(),
}


def verify_response(response: str) -> dict[str, bool]:
    """Check if response found all needles."""
    return {name: check(response) for name, check in NEEDLES.items()}


# =============================================================================
# MAIN
# =============================================================================


def main():
    from headroom.integrations.agno import HeadroomAgnoModel
    from headroom.pricing import estimate_cost

    model_id = "claude-sonnet-4-20250514"

    print("\n" + "=" * 70)
    print("  MULTI-TOOL AGENT DEMO")
    print("=" * 70)

    # Show data sizes
    total_chars = sum(
        len(json.dumps(d)) for d in [GITHUB_ISSUES, CODE_RESULTS, DB_LOGS, ARXIV_PAPERS]
    )
    print(f"\n  Tool outputs: {total_chars:,} chars across 4 tools")
    print("  Needles hidden at positions: #42, #23, #17, #15")

    # Create agent with Headroom
    base_model = Claude(id=model_id)
    model = HeadroomAgnoModel(wrapped_model=base_model)
    agent = Agent(
        model=model,
        tools=[search_github_issues, search_code, query_database, search_arxiv],
        markdown=True,
    )

    question = """Investigate a memory leak in our application:
1. Search GitHub for memory-related issues
2. Search code for memory leak fixes
3. Check database logs for OutOfMemory errors
4. Find relevant research papers

Summarize findings and recommend a fix."""

    print("\n  Running agent...")
    start = time.time()

    response = agent.run(question)
    response_text = response.content if hasattr(response, "content") else str(response)

    duration = time.time() - start

    # Get stats from Headroom
    stats = model.get_savings_summary()
    tokens_before = stats["total_tokens_before"]
    tokens_after = stats["total_tokens_after"]
    tokens_saved = stats["total_tokens_saved"]
    pct_saved = (tokens_saved / tokens_before * 100) if tokens_before > 0 else 0

    # Calculate costs
    cost_before = estimate_cost(model_id, input_tokens=tokens_before)
    cost_after = estimate_cost(model_id, input_tokens=tokens_after)

    # Verify findings
    verification = verify_response(response_text)
    found = sum(verification.values())

    # Results
    print("\n" + "=" * 70)
    print("  RESULTS")
    print("=" * 70)

    print(f"""
                              Without         With
                              Headroom        Headroom        Savings
    ─────────────────────────────────────────────────────────────────
    Input Tokens              {tokens_before:>8,}        {tokens_after:>8,}        {tokens_saved:,} ({pct_saved:.0f}%)""")

    if cost_before and cost_after:
        cost_saved = cost_before - cost_after
        print(
            f"    Input Cost                ${cost_before:.4f}         ${cost_after:.4f}         ${cost_saved:.4f}"
        )

    print(f"""
    Duration                                    {duration:.1f}s
    API Requests                                {stats["total_requests"]}
    Needles Found                               {found}/4
    """)

    for name, found_it in verification.items():
        print(f"      {'✓' if found_it else '✗'} {name}")

    print("\n" + "=" * 70)
    print("  RESPONSE (excerpt)")
    print("=" * 70)
    print(response_text[:1500] + "..." if len(response_text) > 1500 else response_text)

    print("\n" + "=" * 70)
    print(f"  {pct_saved:.0f}% token reduction, {found}/4 needles found")
    print("=" * 70 + "\n")


if __name__ == "__main__":
    main()