From 1599821c3eb35ac1550aa8eaa45b665e0b5817b9 Mon Sep 17 00:00:00 2001 From: chopratejas Date: Fri, 30 Jan 2026 23:10:44 -0800 Subject: [PATCH 1/3] Add dashboard UI and improve proxy observability - Add web dashboard at /dashboard endpoint with real-time stats - Simplify dashboard metrics to user-friendly terms (removed confusing CCR/TOIN terminology) - Track Headroom overhead separately from total latency - Add request logging to Bedrock paths (was missing) - Use package version (__version__) instead of hardcoded "1.0.0" - Add latency min/max tracking in addition to average Dashboard shows: requests, tokens saved, cost saved, overhead, providers breakdown, performance stats, and recent requests table. --- headroom/dashboard/__init__.py | 12 + headroom/dashboard/templates/dashboard.html | 377 ++++++++++++++++++++ headroom/proxy/server.py | 128 ++++++- 3 files changed, 514 insertions(+), 3 deletions(-) create mode 100644 headroom/dashboard/__init__.py create mode 100644 headroom/dashboard/templates/dashboard.html diff --git a/headroom/dashboard/__init__.py b/headroom/dashboard/__init__.py new file mode 100644 index 000000000..74010394d --- /dev/null +++ b/headroom/dashboard/__init__.py @@ -0,0 +1,12 @@ +"""Headroom Dashboard - Real-time proxy monitoring UI.""" + +from pathlib import Path + +DASHBOARD_DIR = Path(__file__).parent +TEMPLATES_DIR = DASHBOARD_DIR / "templates" + + +def get_dashboard_html() -> str: + """Load the dashboard HTML template.""" + template_path = TEMPLATES_DIR / "dashboard.html" + return template_path.read_text() diff --git a/headroom/dashboard/templates/dashboard.html b/headroom/dashboard/templates/dashboard.html new file mode 100644 index 000000000..9e02420db --- /dev/null +++ b/headroom/dashboard/templates/dashboard.html @@ -0,0 +1,377 @@ + + + + + + Headroom Dashboard + + + + + + + + +
+
+

HEADROOM

+ +
+
+
+ Status + + + + +
+
+ Updated +
+
+
+ +
+ +
+ +
+
Requests
+
+
+ + + + + + + + + + +
+
+ + +
+
Tokens Saved
+
+ + +
+
+ + + + +
+
+ + +
+
Cost Saved
+
+ +
+
+ vs $ spent +
+
+ + +
+
Headroom Overhead
+
+ +
+
+ Avg s total response time +
+
+
+ + +
+ +
+
Token Usage
+
+
+ Input Tokens + +
+
+ Output Tokens + +
+
+
+ Original Size + +
+
+ After Compression + +
+
+
+ + +
+
Providers
+
+ + +
+
+ + +
+
Performance
+
+
+ Headroom Overhead + +
+
+ Overhead Range + +
+
+ Total Response Time + +
+
+ Failed Requests + +
+
+
+
+ + +
+
+ Recent Requests + Last 10 +
+
+ + + + + + + + + + + + + + + + +
TimeModelInputOutputSavedCostLatency
+
+
+ + + +
+ + + + + + + diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index c22bb8aa1..abf7b2228 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -44,7 +44,7 @@ try: import uvicorn from fastapi import FastAPI, HTTPException, Request, Response from fastapi.middleware.cors import CORSMiddleware - from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse + from fastapi.responses import HTMLResponse, JSONResponse, PlainTextResponse, StreamingResponse FASTAPI_AVAILABLE = True except ImportError: @@ -53,6 +53,7 @@ except ImportError: # Add parent to path for imports sys.path.insert(0, str(Path(__file__).parent.parent.parent)) +from headroom import __version__ from headroom.backends import LiteLLMBackend from headroom.backends.base import Backend from headroom.cache.compression_feedback import get_compression_feedback @@ -78,6 +79,7 @@ from headroom.config import ( RollingWindowConfig, SmartCrusherConfig, ) +from headroom.dashboard import get_dashboard_html from headroom.providers import AnthropicProvider, OpenAIProvider from headroom.proxy.memory_handler import MemoryConfig, MemoryHandler from headroom.telemetry import get_telemetry_collector @@ -676,8 +678,15 @@ class PrometheusMetrics: self.tokens_saved_total = 0 self.latency_sum_ms = 0.0 + self.latency_min_ms = float("inf") + self.latency_max_ms = 0.0 self.latency_count = 0 + # Headroom overhead (optimization time only, excludes LLM) + self.overhead_sum_ms = 0.0 + self.overhead_min_ms = float("inf") + self.overhead_max_ms = 0.0 + self.cost_total_usd = 0.0 self.savings_total_usd = 0.0 @@ -694,6 +703,7 @@ class PrometheusMetrics: cached: bool = False, cost_usd: float = 0, savings_usd: float = 0, + overhead_ms: float = 0, ): """Record metrics for a request.""" async with self._lock: @@ -709,8 +719,16 @@ class PrometheusMetrics: self.tokens_saved_total += tokens_saved self.latency_sum_ms += latency_ms + self.latency_min_ms = min(self.latency_min_ms, latency_ms) + self.latency_max_ms = max(self.latency_max_ms, latency_ms) self.latency_count += 1 + # Track Headroom overhead separately + if overhead_ms > 0: + self.overhead_sum_ms += overhead_ms + self.overhead_min_ms = min(self.overhead_min_ms, overhead_ms) + self.overhead_max_ms = max(self.overhead_max_ms, overhead_ms) + self.cost_total_usd += cost_usd self.savings_total_usd += savings_usd @@ -1654,14 +1672,51 @@ class HeadroomProxy: tokens_saved=tokens_saved, latency_ms=total_latency, cached=False, + overhead_ms=optimization_latency, ) + cost_usd = None + savings_usd = None if self.cost_tracker: cost_usd = self.cost_tracker.estimate_cost( model, optimized_tokens, output_tokens ) + original_cost = self.cost_tracker.estimate_cost( + model, original_tokens, output_tokens + ) if cost_usd: self.cost_tracker.record_cost(cost_usd) + if cost_usd and original_cost: + savings_usd = original_cost - cost_usd + self.cost_tracker.record_savings(savings_usd) + + # Log request + if self.logger: + self.logger.log( + RequestLog( + request_id=request_id, + timestamp=datetime.now().isoformat(), + provider="bedrock", + model=model, + input_tokens_original=original_tokens, + input_tokens_optimized=optimized_tokens, + output_tokens=output_tokens, + tokens_saved=tokens_saved, + savings_percent=(tokens_saved / original_tokens * 100) + if original_tokens > 0 + else 0, + estimated_cost_usd=cost_usd, + estimated_savings_usd=savings_usd, + optimization_latency_ms=optimization_latency, + total_latency_ms=total_latency, + tags=tags, + cache_hit=False, + transforms_applied=transforms_applied, + request_messages=body.get("messages") + if self.config.log_full_messages + else None, + ) + ) return JSONResponse( status_code=backend_response.status_code, @@ -1909,6 +1964,7 @@ class HeadroomProxy: latency_ms=total_latency, cost_usd=cost_usd or 0, savings_usd=savings_usd or 0, + overhead_ms=optimization_latency, ) # Log request @@ -3676,14 +3732,51 @@ class HeadroomProxy: tokens_saved=tokens_saved, latency_ms=total_latency, cached=False, + overhead_ms=optimization_latency, ) + cost_usd = None + savings_usd = None if self.cost_tracker: cost_usd = self.cost_tracker.estimate_cost( model, optimized_tokens, output_tokens ) + original_cost = self.cost_tracker.estimate_cost( + model, original_tokens, output_tokens + ) if cost_usd: self.cost_tracker.record_cost(cost_usd) + if cost_usd and original_cost: + savings_usd = original_cost - cost_usd + self.cost_tracker.record_savings(savings_usd) + + # Log request + if self.logger: + self.logger.log( + RequestLog( + request_id=request_id, + timestamp=datetime.now().isoformat(), + provider="bedrock", + model=model, + input_tokens_original=original_tokens, + input_tokens_optimized=optimized_tokens, + output_tokens=output_tokens, + tokens_saved=tokens_saved, + savings_percent=(tokens_saved / original_tokens * 100) + if original_tokens > 0 + else 0, + estimated_cost_usd=cost_usd, + estimated_savings_usd=savings_usd, + optimization_latency_ms=optimization_latency, + total_latency_ms=total_latency, + tags=tags, + cache_hit=False, + transforms_applied=transforms_applied, + request_messages=body.get("messages") + if self.config.log_full_messages + else None, + ) + ) if tokens_saved > 0: logger.info( @@ -5231,7 +5324,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: app = FastAPI( title="Headroom Proxy", description="Production-ready LLM optimization proxy", - version="1.0.0", + version=__version__, ) # CORS @@ -5260,7 +5353,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: async def health(): return { "status": "healthy", - "version": "1.0.0", + "version": __version__, "config": { "optimize": config.optimize, "cache": config.cache_enabled, @@ -5268,6 +5361,11 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: }, } + @app.get("/dashboard", response_class=HTMLResponse) + async def dashboard(): + """Serve the Headroom dashboard UI.""" + return get_dashboard_html() + @app.get("/stats") async def stats(): """Get comprehensive proxy statistics. @@ -5284,6 +5382,23 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: # Calculate average latency avg_latency_ms = round(m.latency_sum_ms / m.latency_count, 2) if m.latency_count > 0 else 0 + min_latency_ms = ( + round(m.latency_min_ms, 2) + if m.latency_count > 0 and m.latency_min_ms != float("inf") + else 0 + ) + max_latency_ms = round(m.latency_max_ms, 2) if m.latency_count > 0 else 0 + + # Calculate Headroom overhead (optimization time only) + avg_overhead_ms = ( + round(m.overhead_sum_ms / m.latency_count, 2) if m.latency_count > 0 else 0 + ) + min_overhead_ms = ( + round(m.overhead_min_ms, 2) + if m.latency_count > 0 and m.overhead_min_ms != float("inf") + else 0 + ) + max_overhead_ms = round(m.overhead_max_ms, 2) if m.latency_count > 0 else 0 # Get compression store stats store = get_compression_store() @@ -5323,8 +5438,15 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: }, "latency": { "average_ms": avg_latency_ms, + "min_ms": min_latency_ms, + "max_ms": max_latency_ms, "total_requests": m.latency_count, }, + "overhead": { + "average_ms": avg_overhead_ms, + "min_ms": min_overhead_ms, + "max_ms": max_overhead_ms, + }, "cost": proxy.cost_tracker.stats() if proxy.cost_tracker else None, "compression": { "ccr_entries": compression_stats.get("entry_count", 0), From f3e03f165252fa14f27e8bb8a30d92d5ffa9692e Mon Sep 17 00:00:00 2001 From: chopratejas Date: Fri, 30 Jan 2026 23:12:22 -0800 Subject: [PATCH 2/3] Update README with dashboard info and installation recommendations - Add dashboard URL (http://localhost:8787/dashboard) to quickstart - Recommend headroom-ai[all] for best compression performance - Note that first startup downloads ML models (~500MB one-time) --- README.md | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 1e33b914b..55cc2fdf6 100644 --- a/README.md +++ b/README.md @@ -202,10 +202,14 @@ For deep technical details, see [Architecture Documentation](docs/ARCHITECTURE.m ### Option 1: Proxy (Zero Code Changes) ```bash -pip install "headroom-ai[proxy]" +pip install "headroom-ai[all]" # Recommended for best performance headroom proxy --port 8787 ``` +> **Note:** First startup downloads ML models (~500MB) for optimal compression. This is a one-time download. + +**Dashboard:** Open http://localhost:8787/dashboard to see real-time stats, token savings, and request history. + Point your tools at the proxy: ```bash @@ -434,6 +438,10 @@ New models auto-supported via naming pattern detection. ## Installation ```bash +# Recommended: Install everything for best compression performance +pip install "headroom-ai[all]" + +# Or install specific components pip install headroom-ai # SDK only pip install "headroom-ai[proxy]" # Proxy server pip install "headroom-ai[langchain]" # LangChain integration @@ -441,11 +449,12 @@ pip install "headroom-ai[agno]" # Agno agent framework pip install "headroom-ai[evals]" # Evaluation framework pip install "headroom-ai[code]" # AST-based code compression pip install "headroom-ai[llmlingua]" # ML-based compression -pip install "headroom-ai[all]" # Everything ``` **Requirements**: Python 3.10+ +> **First-time startup:** Headroom downloads ML models (~500MB) on first run for optimal compression. This is cached locally and only happens once. + --- ## Documentation From 95e9b39b6d6f4fb0439a9cb75fcab749e711e523 Mon Sep 17 00:00:00 2001 From: chopratejas Date: Sat, 31 Jan 2026 00:31:37 -0800 Subject: [PATCH 3/3] feat: Add AWS Strands Agents SDK integration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Add Headroom integration with AWS Strands Agents SDK, enabling automatic context optimization and tool output compression for Strands-based agents. Fixes #14 ## Type of Change - [x] New feature (non-breaking change that adds functionality) - [x] Documentation update ## Changes Made ### Core Integration (`headroom/integrations/strands/`) - **HeadroomHookProvider** - Implements Strands `HookProvider` interface for automatic tool output compression via `AfterToolCallEvent`. Compresses verbose tool outputs before they enter conversation context. - **HeadroomStrandsModel** - Model wrapper that extends Strands `Model` base class for message-level optimization. Implements all required abstract methods: `stream()`, `get_config()`, `update_config()`, `structured_output()`. - **Provider auto-detection** - Automatically detects appropriate Headroom provider (Anthropic, OpenAI, Google) based on wrapped Strands model type. - **`strands-agents` as optional dependency** - Install with `pip install headroom-ai[strands]` ### Testing (`tests/integrations/test_strands/`) - **Real integration tests (25 tests)** - Use actual AWS Bedrock API calls with Claude 3 Haiku. Skip automatically when credentials unavailable. - **Unit tests (57 tests)** - Mock-based tests for internal logic, edge cases, and error handling. No credentials required. ### Demo (`examples/strands_bedrock_demo.py`) - Interactive demo showcasing both integration patterns - Visual before/after compression comparison with token savings - 4 verbose tools (search, logs, database, metrics) demonstrating real savings - Supports `--hook` and `--model` flags for individual demos ## Testing All tests verified: - [x] Unit tests pass (57 tests) - [x] Integration tests pass (25 tests with real Bedrock API) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom/integrations/strands/`) - [x] Formatting passes (`ruff format --check`) - [x] Demo runs successfully with ~50% token savings ## Test Output ``` $ pytest tests/integrations/test_strands/ -v =================== 82 passed in 90.09s =================== $ ruff check headroom/integrations/strands/ --ignore E402 All checks passed! $ mypy headroom/integrations/strands/ --ignore-missing-imports Success: no issues found ``` ## Demo Results ``` ╭────────────────────────────────────────────────────────────╮ │ HeadroomHookProvider Results │ │────────────────────────────────────────────────────────────│ │ Tokens BEFORE compression: 51,961 │ │ Tokens AFTER compression: 25,658 │ │ Tokens SAVED: 26,303 (50.6%) │ ╰────────────────────────────────────────────────────────────╯ ``` --- examples/README.md | 56 + examples/strands_bedrock_demo.py | 1001 +++++++++++++++++ headroom/integrations/strands/__init__.py | 88 ++ headroom/integrations/strands/hooks.py | 540 +++++++++ headroom/integrations/strands/model.py | 625 ++++++++++ headroom/integrations/strands/providers.py | 166 +++ pyproject.toml | 4 + tests/integrations/test_strands/__init__.py | 1 + tests/integrations/test_strands/test_hooks.py | 545 +++++++++ .../test_strands/test_hooks_unit.py | 564 ++++++++++ tests/integrations/test_strands/test_model.py | 673 +++++++++++ .../test_strands/test_model_unit.py | 645 +++++++++++ 12 files changed, 4908 insertions(+) create mode 100644 examples/strands_bedrock_demo.py create mode 100644 headroom/integrations/strands/__init__.py create mode 100644 headroom/integrations/strands/hooks.py create mode 100644 headroom/integrations/strands/model.py create mode 100644 headroom/integrations/strands/providers.py create mode 100644 tests/integrations/test_strands/__init__.py create mode 100644 tests/integrations/test_strands/test_hooks.py create mode 100644 tests/integrations/test_strands/test_hooks_unit.py create mode 100644 tests/integrations/test_strands/test_model.py create mode 100644 tests/integrations/test_strands/test_model_unit.py diff --git a/examples/README.md b/examples/README.md index e4fbf240d..117a5849c 100644 --- a/examples/README.md +++ b/examples/README.md @@ -86,6 +86,44 @@ export OPENAI_API_KEY='your-key' PYTHONPATH=. python -m examples.mcp_demo.run_agent_eval ``` +### strands_bedrock_demo.py + +AWS Strands Agents + Bedrock integration demo. Showcases two Headroom integration patterns: + +1. **HeadroomHookProvider** - Compresses tool outputs in real-time +2. **HeadroomStrandsModel** - Optimizes entire conversation context + +```bash +# Configure AWS credentials +export AWS_ACCESS_KEY_ID='your-access-key' +export AWS_SECRET_ACCESS_KEY='your-secret-key' +export AWS_DEFAULT_REGION='us-west-2' # Optional, defaults to us-west-2 + +# Or use AWS profile +export AWS_PROFILE='your-profile-name' + +# Run the full demo (both integration patterns) +python examples/strands_bedrock_demo.py + +# Run only the hook provider demo +python examples/strands_bedrock_demo.py --hook + +# Run only the model wrapper demo +python examples/strands_bedrock_demo.py --model + +# Specify a different AWS region +python examples/strands_bedrock_demo.py --region us-east-1 +``` + +The demo uses Claude 3 Haiku via Bedrock for cost efficiency. It creates agents with +4 tools that return verbose JSON output (search results, logs, database records, metrics) +and displays compression statistics with visual comparisons. + +**Requirements:** +- AWS account with Bedrock enabled +- Claude 3 Haiku model access in your region +- `pip install strands-agents headroom-ai[strands]` + ## Running Examples All examples can be run from the repository root: @@ -105,6 +143,7 @@ python examples/.py | basic_usage | 50-70% | Simple tool output compression | | langchain_demo | 70-85% | Real agent with multiple tools | | mcp_demo | 60-80% | MCP tool outputs | +| strands_bedrock_demo | 60-85% | Strands + Bedrock with verbose tools | | real_world_eval | 50-90% | Varies by scenario | ## Troubleshooting @@ -131,3 +170,20 @@ Ensure your API keys are set: export OPENAI_API_KEY='sk-...' export ANTHROPIC_API_KEY='sk-ant-...' ``` + +**AWS Credentials Errors (for Strands demo)** + +Ensure AWS credentials are configured: + +```bash +# Option 1: Environment variables +export AWS_ACCESS_KEY_ID='your-access-key' +export AWS_SECRET_ACCESS_KEY='your-secret-key' + +# Option 2: AWS profile +export AWS_PROFILE='your-profile-name' + +# Option 3: AWS credentials file (~/.aws/credentials) +``` + +Also ensure Bedrock and the Claude 3 Haiku model are enabled in your AWS account. diff --git a/examples/strands_bedrock_demo.py b/examples/strands_bedrock_demo.py new file mode 100644 index 000000000..74f5c1de9 --- /dev/null +++ b/examples/strands_bedrock_demo.py @@ -0,0 +1,1001 @@ +#!/usr/bin/env python3 +"""Comprehensive Strands + Bedrock Demo for Headroom SDK. + +This demo showcases two Headroom integration patterns for AWS Strands Agents: + +1. **HeadroomHookProvider** - Compresses tool outputs as they happen + - Intercepts tool results via Strands hooks + - Applies SmartCrusher compression to large JSON outputs + - Shows per-tool compression metrics + +2. **HeadroomStrandsModel** - Optimizes entire conversation context + - Wraps BedrockModel for automatic context optimization + - Applies message-level transforms before API calls + - Tracks cumulative savings across the session + +Run with: + python examples/strands_bedrock_demo.py # Run both demos + python examples/strands_bedrock_demo.py --hook # Hook provider demo only + python examples/strands_bedrock_demo.py --model # Model wrapper demo only + +Requirements: + - AWS credentials configured (AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY or AWS_PROFILE) + - pip install strands-agents headroom-ai[strands] +""" + +from __future__ import annotations + +import argparse +import json +import os +import random +import sys +from datetime import datetime, timedelta +from typing import Any + +# ============================================================================ +# Check Dependencies +# ============================================================================ + + +def check_dependencies() -> bool: + """Check if required dependencies are available.""" + missing = [] + + # Check strands-agents + try: + from strands import Agent # noqa: F401 + from strands.models import BedrockModel # noqa: F401 + except ImportError: + missing.append("strands-agents") + + # Check headroom + try: + from headroom.integrations.strands import ( # noqa: F401 + HeadroomHookProvider, + HeadroomStrandsModel, + ) + except ImportError: + missing.append("headroom-ai[strands]") + + if missing: + print_box( + "Missing Dependencies", + [ + "The following packages are required but not installed:", + "", + *[f" - {pkg}" for pkg in missing], + "", + "Install with:", + f" pip install {' '.join(missing)}", + ], + style="error", + ) + return False + + return True + + +def check_aws_credentials() -> bool: + """Check if AWS credentials are available.""" + has_env_keys = os.environ.get("AWS_ACCESS_KEY_ID") and os.environ.get("AWS_SECRET_ACCESS_KEY") + has_profile = os.environ.get("AWS_PROFILE") + has_creds_file = os.path.exists(os.path.expanduser("~/.aws/credentials")) + + if not (has_env_keys or has_profile or has_creds_file): + print_box( + "AWS Credentials Not Found", + [ + "This demo requires AWS credentials to access Bedrock.", + "", + "Configure credentials using one of these methods:", + "", + "1. Environment variables:", + " export AWS_ACCESS_KEY_ID='your-access-key'", + " export AWS_SECRET_ACCESS_KEY='your-secret-key'", + " export AWS_DEFAULT_REGION='us-west-2'", + "", + "2. AWS Profile:", + " export AWS_PROFILE='your-profile-name'", + "", + "3. AWS credentials file:", + " ~/.aws/credentials", + ], + style="error", + ) + return False + + return True + + +# ============================================================================ +# Pretty Printing Utilities +# ============================================================================ + + +def print_box(title: str, lines: list[str], style: str = "normal", width: int = 76) -> None: + """Print a box with title and content using box drawing characters.""" + if style == "error": + top_left, top_right = "\u2554", "\u2557" # Double line + bot_left, bot_right = "\u255a", "\u255d" + horiz, vert = "\u2550", "\u2551" + elif style == "success": + top_left, top_right = "\u256d", "\u256e" # Rounded + bot_left, bot_right = "\u2570", "\u256f" + horiz, vert = "\u2500", "\u2502" + else: + top_left, top_right = "\u250c", "\u2510" # Normal single + bot_left, bot_right = "\u2514", "\u2518" + horiz, vert = "\u2500", "\u2502" + + print() + print(f"{top_left}{horiz * (width - 2)}{top_right}") + + # Title + title_padding = (width - 4 - len(title)) // 2 + print( + f"{vert} {' ' * title_padding}{title}{' ' * (width - 4 - title_padding - len(title))} {vert}" + ) + print(f"{vert}{horiz * (width - 2)}{vert}") + + # Content lines + for line in lines: + # Handle lines longer than width + if len(line) > width - 4: + line = line[: width - 7] + "..." + padding = width - 4 - len(line) + print(f"{vert} {line}{' ' * padding} {vert}") + + print(f"{bot_left}{horiz * (width - 2)}{bot_right}") + print() + + +def print_metrics_table( + metrics: list[dict[str, Any]], + headers: list[str], + keys: list[str], + title: str = "Metrics", +) -> None: + """Print metrics in a formatted table.""" + # Calculate column widths + col_widths = [] + for i, header in enumerate(headers): + max_width = len(header) + for m in metrics: + val = m.get(keys[i], "") + max_width = max(max_width, len(str(val))) + col_widths.append(min(max_width + 2, 25)) + + total_width = sum(col_widths) + len(col_widths) + 1 + + print(f"\n {title}") + print(" " + "\u2500" * (total_width - 2)) + + # Header row + header_row = "\u2502" + for i, header in enumerate(headers): + header_row += f" {header:<{col_widths[i] - 2}} \u2502" + print(" " + header_row) + print(" " + "\u2502" + "\u2500" * (total_width - 2) + "\u2502") + + # Data rows + for m in metrics: + row = "\u2502" + for i, key in enumerate(keys): + val = str(m.get(key, "")) + if len(val) > col_widths[i] - 2: + val = val[: col_widths[i] - 5] + "..." + row += f" {val:<{col_widths[i] - 2}} \u2502" + print(" " + row) + + print(" " + "\u2500" * total_width) + + +def print_comparison(before: int, after: int, label: str = "Tokens") -> None: + """Print a before/after comparison with savings.""" + saved = before - after + pct = (saved / before * 100) if before > 0 else 0 + + bar_width = 40 + before_bar = int((before / max(before, 1)) * bar_width) + after_bar = int((after / max(before, 1)) * bar_width) + + print(f"\n {label} Comparison:") + print(f" BEFORE: {before:>8,} \u2502{'=' * before_bar}") + print(f" AFTER: {after:>8,} \u2502{'=' * after_bar}") + print(f" SAVED: {saved:>8,} ({pct:.1f}%)") + + +# ============================================================================ +# Mock Tools - Generate Verbose Output +# ============================================================================ + + +def search_documentation(query: str, limit: int = 25) -> str: + """Search documentation for matching articles. + + Returns search results with titles, snippets, URLs, and metadata. + Simulates a real documentation search API returning verbose results. + """ + results = [] + categories = [ + "getting-started", + "api-reference", + "tutorials", + "troubleshooting", + "best-practices", + ] + sources = ["internal-docs", "confluence", "notion", "github-wiki", "readme"] + + for i in range(limit): + result = { + "id": f"doc-{random.randint(10000, 99999)}", + "title": f"{query.title()} Guide - Part {i + 1}", + "snippet": f"This comprehensive guide covers {query} implementation. " + f"Learn how to configure, deploy, and maintain {query} in production. " + f"Includes examples, best practices, and troubleshooting tips for {query}.", + "url": f"https://docs.example.com/{query.replace(' ', '-')}/section-{i + 1}", + "category": random.choice(categories), + "source": random.choice(sources), + "relevance_score": round(random.uniform(0.5, 1.0), 3), + "last_updated": (datetime.now() - timedelta(days=random.randint(1, 180))).isoformat(), + "author": f"Author {random.randint(1, 20)}", + "word_count": random.randint(500, 5000), + "views": random.randint(100, 10000), + "helpful_votes": random.randint(10, 500), + "tags": random.sample( + ["aws", "python", "deployment", "security", "performance", "monitoring"], + k=random.randint(2, 4), + ), + } + results.append(result) + + results.sort(key=lambda x: x["relevance_score"], reverse=True) + + return json.dumps( + { + "query": query, + "total_results": limit * 5, # Simulate more results available + "page": 1, + "per_page": limit, + "results": results, + }, + indent=2, + ) + + +def get_server_logs(server: str, lines: int = 100) -> str: + """Fetch server logs for analysis. + + Returns JSON log entries with timestamps, levels, messages, and context. + Simulates verbose application logs with mostly INFO entries and some errors. + """ + entries = [] + levels = ["DEBUG", "INFO", "INFO", "INFO", "INFO", "WARN", "ERROR"] + services = ["api-gateway", "auth-service", "data-processor", "cache-layer", "message-queue"] + + for _i in range(lines): + timestamp = datetime.now() - timedelta(minutes=random.randint(1, 1440)) + level = random.choice(levels) + + if level == "ERROR": + message = random.choice( + [ + f"Connection timeout to {server}-db after 30000ms", + "Failed to authenticate request: invalid JWT signature", + "Rate limit exceeded for client IP 10.0.0.42", + "Database query failed: connection pool exhausted", + f"Service {server} health check failed: connection refused", + ] + ) + elif level == "WARN": + message = random.choice( + [ + f"Slow query detected on {server}: execution time 2.5s", + "Memory usage at 85% - consider scaling", + "Retry attempt 2/3 for downstream service call", + "Certificate expires in 7 days - renewal required", + ] + ) + else: + message = f"Request processed successfully - endpoint=/api/v1/{server}/data" + + entry = { + "timestamp": timestamp.isoformat(), + "level": level, + "server": server, + "service": random.choice(services), + "message": message, + "trace_id": f"trace-{random.randint(100000, 999999):06x}", + "span_id": f"span-{random.randint(1000, 9999):04x}", + "request_id": f"req-{random.randint(10000000, 99999999)}", + "client_ip": f"10.0.{random.randint(0, 255)}.{random.randint(1, 254)}", + "user_agent": random.choice( + [ + "Mozilla/5.0 (compatible; MonitorBot/1.0)", + "python-requests/2.31.0", + "curl/8.1.2", + "PostmanRuntime/7.32.0", + ] + ), + "response_time_ms": random.randint(5, 2000), + "status_code": 200 + if level in ["DEBUG", "INFO"] + else random.choice([400, 500, 502, 503]), + "metadata": { + "pod": f"{server}-{random.randint(1, 5)}-abc123", + "node": f"ip-10-0-{random.randint(0, 255)}-{random.randint(1, 254)}.ec2.internal", + "region": random.choice(["us-west-2", "us-east-1", "eu-west-1"]), + "version": f"v1.{random.randint(0, 9)}.{random.randint(0, 20)}", + }, + } + entries.append(entry) + + entries.sort(key=lambda x: x["timestamp"], reverse=True) + + return json.dumps( + { + "server": server, + "log_count": lines, + "time_range": { + "start": entries[-1]["timestamp"] if entries else None, + "end": entries[0]["timestamp"] if entries else None, + }, + "entries": entries, + }, + indent=2, + ) + + +def query_database(sql: str, limit: int = 50) -> str: + """Execute a database query and return results. + + Returns rows of data as if from a real database query. + Simulates customer/order/transaction data. + """ + # Parse table name from SQL (simple simulation) + table = "records" + for word in sql.lower().split(): + if word in ["users", "orders", "transactions", "customers", "products", "events"]: + table = word + break + + rows = [] + statuses = ["active", "pending", "completed", "cancelled", "refunded"] + + for i in range(limit): + if table == "users": + row = { + "user_id": f"usr-{random.randint(100000, 999999)}", + "email": f"user{i}@example.com", + "name": f"Customer {i}", + "status": random.choice(["active", "inactive", "suspended"]), + "created_at": (datetime.now() - timedelta(days=random.randint(1, 365))).isoformat(), + "last_login": ( + datetime.now() - timedelta(hours=random.randint(1, 720)) + ).isoformat(), + "plan": random.choice(["free", "basic", "pro", "enterprise"]), + "country": random.choice(["US", "UK", "DE", "FR", "JP", "AU"]), + } + elif table == "orders": + row = { + "order_id": f"ord-{random.randint(100000, 999999)}", + "customer_id": f"usr-{random.randint(100000, 999999)}", + "total": round(random.uniform(10, 1000), 2), + "currency": random.choice(["USD", "EUR", "GBP"]), + "status": random.choice(statuses), + "items_count": random.randint(1, 10), + "created_at": (datetime.now() - timedelta(days=random.randint(1, 90))).isoformat(), + "shipped_at": (datetime.now() - timedelta(days=random.randint(0, 30))).isoformat() + if random.random() > 0.3 + else None, + } + else: + row = { + "id": i + 1, + "record_type": table, + "value": random.randint(100, 10000), + "status": random.choice(statuses), + "created_at": (datetime.now() - timedelta(days=random.randint(1, 365))).isoformat(), + "metadata": { + "source": random.choice(["web", "api", "import", "sync"]), + "version": f"v{random.randint(1, 5)}", + }, + } + rows.append(row) + + return json.dumps( + { + "query": sql, + "table": table, + "row_count": limit, + "total_available": limit * 10, + "execution_time_ms": random.randint(10, 500), + "rows": rows, + }, + indent=2, + ) + + +def get_system_metrics(timerange: str = "1h", service: str = "all") -> str: + """Get system metrics for monitoring. + + Returns time-series data points for CPU, memory, latency, and error rates. + Simulates Prometheus/CloudWatch style metrics. + """ + # Parse timerange to determine number of points + points = {"5m": 10, "15m": 30, "1h": 60, "6h": 72, "24h": 144}.get(timerange, 60) + + data_points = [] + services_list = ["api", "worker", "cache", "database"] if service == "all" else [service] + + for svc in services_list: + for i in range(points): + timestamp = datetime.now() - timedelta(minutes=i * (60 // min(points, 60))) + + # Inject some anomalies + is_anomaly = random.random() < 0.05 + + point = { + "timestamp": timestamp.isoformat(), + "service": svc, + "metrics": { + "cpu_percent": round( + random.uniform(70, 95) if is_anomaly else random.uniform(20, 45), 2 + ), + "memory_percent": round( + random.uniform(80, 95) if is_anomaly else random.uniform(40, 65), 2 + ), + "memory_mb": random.randint(2000, 4000) + if is_anomaly + else random.randint(500, 1500), + "latency_p50_ms": random.randint(100, 500) + if is_anomaly + else random.randint(10, 50), + "latency_p95_ms": random.randint(500, 2000) + if is_anomaly + else random.randint(50, 150), + "latency_p99_ms": random.randint(1000, 5000) + if is_anomaly + else random.randint(100, 300), + "request_rate_per_sec": random.randint(500, 2000) + if is_anomaly + else random.randint(50, 200), + "error_rate_percent": round( + random.uniform(5, 15) if is_anomaly else random.uniform(0, 1), 3 + ), + "active_connections": random.randint(200, 500) + if is_anomaly + else random.randint(20, 80), + }, + "health": "degraded" if is_anomaly else "healthy", + "region": random.choice(["us-west-2", "us-east-1", "eu-west-1"]), + } + data_points.append(point) + + # Calculate summary statistics + all_cpu = [p["metrics"]["cpu_percent"] for p in data_points] + all_mem = [p["metrics"]["memory_percent"] for p in data_points] + all_latency = [p["metrics"]["latency_p50_ms"] for p in data_points] + + return json.dumps( + { + "timerange": timerange, + "service": service, + "data_points_count": len(data_points), + "summary": { + "cpu": { + "min": min(all_cpu), + "max": max(all_cpu), + "avg": sum(all_cpu) / len(all_cpu), + }, + "memory": { + "min": min(all_mem), + "max": max(all_mem), + "avg": sum(all_mem) / len(all_mem), + }, + "latency_p50": { + "min": min(all_latency), + "max": max(all_latency), + "avg": sum(all_latency) / len(all_latency), + }, + }, + "data_points": data_points, + }, + indent=2, + ) + + +# ============================================================================ +# Demo 1: HeadroomHookProvider +# ============================================================================ + + +def run_hook_provider_demo(region: str = "us-west-2") -> dict[str, Any]: + """Demonstrate HeadroomHookProvider for tool output compression. + + Returns metrics from the demo run. + """ + from strands import Agent, tool + from strands.models import BedrockModel + + from headroom.integrations.strands import HeadroomHookProvider + + print_box( + "Demo 1: HeadroomHookProvider", + [ + "The HeadroomHookProvider intercepts tool outputs and compresses", + "them BEFORE they're added to the conversation context.", + "", + "This reduces token usage for subsequent LLM calls by eliminating", + "redundant data from verbose tool outputs.", + "", + "Using: Claude 3 Haiku (anthropic.claude-3-haiku-20240307-v1:0)", + ], + ) + + # Define tools with @tool decorator + @tool + def search_docs_tool(query: str) -> str: + """Search documentation for articles matching the query. + + Args: + query: The search query to find relevant documentation + + Returns: + JSON array of search results with titles, snippets, and URLs + """ + return search_documentation(query, limit=25) + + @tool + def get_logs_tool(server: str, lines: int = 100) -> str: + """Fetch server logs for analysis and troubleshooting. + + Args: + server: Name of the server to fetch logs from + lines: Number of log lines to retrieve (default: 100) + + Returns: + JSON array of log entries with timestamps and messages + """ + return get_server_logs(server, lines=lines) + + @tool + def query_db_tool(sql: str) -> str: + """Execute a database query and return results. + + Args: + sql: SQL query to execute (e.g., SELECT * FROM users) + + Returns: + JSON array of database rows + """ + return query_database(sql, limit=50) + + @tool + def get_metrics_tool(timerange: str = "1h") -> str: + """Get system metrics for the specified time range. + + Args: + timerange: Time range for metrics (5m, 15m, 1h, 6h, 24h) + + Returns: + JSON object with time-series metrics data + """ + return get_system_metrics(timerange) + + # Create BedrockModel + model = BedrockModel( + model_id="anthropic.claude-3-haiku-20240307-v1:0", + region_name=region, + temperature=0.1, + ) + + # Create HeadroomHookProvider + hook_provider = HeadroomHookProvider( + compress_tool_outputs=True, + min_tokens_to_compress=100, # Compress outputs with 100+ tokens + preserve_errors=True, + ) + + # Create agent with hook + agent = Agent( + model=model, + tools=[search_docs_tool, get_logs_tool, query_db_tool, get_metrics_tool], + hooks=[hook_provider], + ) + + print("\n Running agent queries that trigger tools with verbose output...") + print(" " + "-" * 60) + + # Query 1: Search documentation + print("\n Query 1: Searching documentation...") + result1 = agent( + "Search the documentation for 'authentication setup' and summarize " + "the top 3 most relevant articles you find." + ) + print(f" Response: {str(result1)[:200]}...") + + # Query 2: Get server logs + print("\n Query 2: Fetching server logs...") + result2 = agent( + "Get the logs from server 'api-gateway' (100 lines) and tell me " + "how many ERROR and WARN level entries there are." + ) + print(f" Response: {str(result2)[:200]}...") + + # Query 3: Query database + print("\n Query 3: Running database query...") + result3 = agent( + "Query the orders table and tell me how many orders have status 'completed' " + "and what the average order total is." + ) + print(f" Response: {str(result3)[:200]}...") + + # Query 4: Get metrics + print("\n Query 4: Fetching system metrics...") + result4 = agent( + "Get the system metrics for the last hour and identify if there are " + "any services with high CPU usage (>70%) or memory issues." + ) + print(f" Response: {str(result4)[:200]}...") + + # Get metrics + metrics = hook_provider.get_savings_summary() + + # Display results + print_box( + "HeadroomHookProvider Results", + [ + f"Tool calls processed: {metrics['total_requests']}", + f"Compressions applied: {metrics['compressed_requests']}", + "", + f"Tokens BEFORE compression: {metrics['total_tokens_before']:,}", + f"Tokens AFTER compression: {metrics['total_tokens_after']:,}", + f"Tokens SAVED: {metrics['total_tokens_saved']:,}", + "", + f"Average savings: {metrics['average_savings_percent']:.1f}%", + ], + style="success", + ) + + # Show per-tool breakdown + if hook_provider.metrics_history: + tool_metrics = [] + for m in hook_provider.metrics_history: + tool_metrics.append( + { + "tool": m.tool_name[:20], + "before": f"{m.tokens_before:,}", + "after": f"{m.tokens_after:,}", + "saved": f"{m.tokens_saved:,}", + "pct": f"{m.savings_percent:.1f}%", + } + ) + + print_metrics_table( + tool_metrics, + headers=["Tool", "Before", "After", "Saved", "%"], + keys=["tool", "before", "after", "saved", "pct"], + title="Per-Tool Compression Breakdown", + ) + + print_comparison( + metrics["total_tokens_before"], + metrics["total_tokens_after"], + "Tool Output Tokens", + ) + + return metrics + + +# ============================================================================ +# Demo 2: HeadroomStrandsModel +# ============================================================================ + + +def run_model_wrapper_demo(region: str = "us-west-2") -> dict[str, Any]: + """Demonstrate HeadroomStrandsModel for conversation optimization. + + Returns metrics from the demo run. + """ + from strands import Agent, tool + from strands.models import BedrockModel + + from headroom import HeadroomConfig + from headroom.integrations.strands import HeadroomStrandsModel + + print_box( + "Demo 2: HeadroomStrandsModel", + [ + "HeadroomStrandsModel wraps the Bedrock model to optimize the", + "ENTIRE conversation context before each API call.", + "", + "As conversations grow with tool outputs and history, the", + "model wrapper applies transforms to reduce context size.", + "", + "Using: Claude 3 Haiku wrapped with HeadroomStrandsModel", + ], + ) + + # Define tools + @tool + def verbose_search(query: str) -> str: + """Search for information with verbose results. + + Args: + query: Search query + + Returns: + Detailed search results + """ + return search_documentation(query, limit=30) + + @tool + def verbose_logs(server: str) -> str: + """Get verbose server logs. + + Args: + server: Server name + + Returns: + Detailed log entries + """ + return get_server_logs(server, lines=150) + + @tool + def verbose_metrics(timerange: str = "1h") -> str: + """Get verbose metrics data. + + Args: + timerange: Time range + + Returns: + Detailed metrics + """ + return get_system_metrics(timerange) + + @tool + def verbose_database(table: str) -> str: + """Query database with verbose results. + + Args: + table: Table name to query + + Returns: + Database records + """ + return query_database(f"SELECT * FROM {table}", limit=60) + + # Create base Bedrock model + base_model = BedrockModel( + model_id="anthropic.claude-3-haiku-20240307-v1:0", + region_name=region, + temperature=0.1, + ) + + # Configure Headroom + config = HeadroomConfig() + config.smart_crusher.enabled = True + config.smart_crusher.min_tokens_to_crush = 100 + config.smart_crusher.max_items_after_crush = 20 + + # Wrap with HeadroomStrandsModel + optimized_model = HeadroomStrandsModel( + wrapped_model=base_model, + config=config, + auto_detect_provider=True, + ) + + # Create agent + agent = Agent( + model=optimized_model, + tools=[verbose_search, verbose_logs, verbose_metrics, verbose_database], + ) + + print("\n Building up a multi-turn conversation with verbose tool outputs...") + print(" " + "-" * 60) + + # Simulate a multi-turn conversation + turns = [ + ("Turn 1", "Search for documentation about 'kubernetes deployment' and give me a summary."), + ("Turn 2", "Now get the logs from the 'worker-service' server and identify any errors."), + ("Turn 3", "Query the orders database and tell me the distribution of order statuses."), + ("Turn 4", "Get the system metrics for the last hour and highlight any anomalies."), + ("Turn 5", "Based on everything you've found, what's the overall system health status?"), + ] + + for turn_name, query in turns: + print(f"\n {turn_name}: {query[:60]}...") + result = agent(query) + print(f" Response: {str(result)[:150]}...") + + # Get metrics + metrics = optimized_model.get_savings_summary() + + # Display results + print_box( + "HeadroomStrandsModel Results", + [ + f"API calls made: {metrics['total_requests']}", + "", + f"Total tokens BEFORE opt: {metrics['total_tokens_before']:,}", + f"Total tokens AFTER opt: {metrics['total_tokens_after']:,}", + f"Total tokens SAVED: {metrics['total_tokens_saved']:,}", + "", + f"Average savings per call: {metrics['average_savings_percent']:.1f}%", + ], + style="success", + ) + + # Show per-request breakdown + if optimized_model.metrics_history: + request_metrics = [] + for i, m in enumerate(optimized_model.metrics_history): + request_metrics.append( + { + "request": f"Request {i + 1}", + "before": f"{m.tokens_before:,}", + "after": f"{m.tokens_after:,}", + "saved": f"{m.tokens_saved:,}", + "pct": f"{m.savings_percent:.1f}%", + } + ) + + print_metrics_table( + request_metrics, + headers=["Request", "Before", "After", "Saved", "%"], + keys=["request", "before", "after", "saved", "pct"], + title="Per-Request Optimization Breakdown", + ) + + print_comparison( + metrics["total_tokens_before"], + metrics["total_tokens_after"], + "Conversation Tokens", + ) + + return metrics + + +# ============================================================================ +# Main +# ============================================================================ + + +def main() -> int: + """Run the Strands Bedrock demo.""" + parser = argparse.ArgumentParser( + description="Headroom + Strands Bedrock Demo", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + python examples/strands_bedrock_demo.py # Run both demos + python examples/strands_bedrock_demo.py --hook # Hook provider only + python examples/strands_bedrock_demo.py --model # Model wrapper only + +Environment Variables: + AWS_ACCESS_KEY_ID AWS access key + AWS_SECRET_ACCESS_KEY AWS secret key + AWS_DEFAULT_REGION AWS region (default: us-west-2) + AWS_PROFILE AWS profile name (alternative to keys) + """, + ) + parser.add_argument( + "--hook", + action="store_true", + help="Run only the HeadroomHookProvider demo", + ) + parser.add_argument( + "--model", + action="store_true", + help="Run only the HeadroomStrandsModel demo", + ) + parser.add_argument( + "--region", + default=os.environ.get("AWS_DEFAULT_REGION", "us-west-2"), + help="AWS region for Bedrock (default: us-west-2)", + ) + + args = parser.parse_args() + + # If neither flag is set, run both + run_hook = args.hook or (not args.hook and not args.model) + run_model = args.model or (not args.hook and not args.model) + + # Print header + print_box( + "Headroom + Strands Bedrock Demo", + [ + "This demo showcases Headroom's integration with AWS Strands Agents.", + "", + "Headroom provides two integration patterns:", + " 1. HeadroomHookProvider - Compress tool outputs in real-time", + " 2. HeadroomStrandsModel - Optimize entire conversation context", + "", + f"Region: {args.region}", + "Model: Claude 3 Haiku (fast and cost-effective for demos)", + ], + ) + + # Check dependencies + if not check_dependencies(): + return 1 + + # Check AWS credentials + if not check_aws_credentials(): + return 1 + + print("\n All checks passed. Starting demos...\n") + + all_metrics = {} + + try: + # Run hook provider demo + if run_hook: + hook_metrics = run_hook_provider_demo(region=args.region) + all_metrics["hook_provider"] = hook_metrics + + # Run model wrapper demo + if run_model: + model_metrics = run_model_wrapper_demo(region=args.region) + all_metrics["model_wrapper"] = model_metrics + + # Print final summary + if run_hook and run_model: + total_before = all_metrics.get("hook_provider", {}).get( + "total_tokens_before", 0 + ) + all_metrics.get("model_wrapper", {}).get("total_tokens_before", 0) + total_after = all_metrics.get("hook_provider", {}).get( + "total_tokens_after", 0 + ) + all_metrics.get("model_wrapper", {}).get("total_tokens_after", 0) + total_saved = total_before - total_after + total_pct = (total_saved / total_before * 100) if total_before > 0 else 0 + + # Estimate cost savings (Claude 3 Haiku pricing) + # Input: $0.25 / 1M tokens, Output: $1.25 / 1M tokens + cost_per_token = 0.25 / 1_000_000 + cost_saved = total_saved * cost_per_token + + print_box( + "Session Summary", + [ + "Combined metrics from both demos:", + "", + f"Total tokens processed: {total_before:,}", + f"Total tokens after opt: {total_after:,}", + f"Total tokens saved: {total_saved:,} ({total_pct:.1f}%)", + "", + f"Estimated cost savings: ${cost_saved:.6f}", + "(At scale, these savings compound significantly!)", + "", + "Integration patterns demonstrated:", + " [x] HeadroomHookProvider - Real-time tool output compression", + " [x] HeadroomStrandsModel - Full context optimization", + ], + style="success", + ) + + return 0 + + except Exception as e: + print_box( + "Error", + [ + f"An error occurred: {type(e).__name__}", + "", + str(e)[:200], + "", + "Common issues:", + " - Invalid AWS credentials", + " - Bedrock not enabled in your AWS account", + " - Model not available in selected region", + " - Rate limiting from too many requests", + ], + style="error", + ) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/headroom/integrations/strands/__init__.py b/headroom/integrations/strands/__init__.py new file mode 100644 index 000000000..a3a1ef0cf --- /dev/null +++ b/headroom/integrations/strands/__init__.py @@ -0,0 +1,88 @@ +"""Strands Agents integration for Headroom SDK. + +This module provides seamless integration with Strands Agents, +enabling automatic context optimization for Strands agents. + +Components: +1. HeadroomStrandsModel - Wraps any Strands model to apply Headroom transforms +2. HeadroomHookProvider - Hook provider for Strands agents +3. get_headroom_provider - Detects appropriate provider for a Strands model +4. get_model_name_from_strands - Extracts model name from a Strands model + +Example: + from strands import Agent + from strands.models import BedrockModel + from headroom.integrations.strands import HeadroomStrandsModel + + # Wrap any Strands model + model = BedrockModel(model_id="anthropic.claude-3-5-sonnet-20241022-v2:0") + optimized_model = HeadroomStrandsModel(model) + + # Use with agent + agent = Agent(model=optimized_model) + response = agent("Hello!") +""" + +from __future__ import annotations + +import importlib.util +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .hooks import HeadroomHookProvider + from .model import HeadroomStrandsModel, OptimizationMetrics, optimize_messages + from .providers import get_headroom_provider, get_model_name_from_strands + + +def strands_available() -> bool: + """Check if strands-agents is installed and available. + + Returns: + True if strands-agents package is available, False otherwise. + """ + return importlib.util.find_spec("strands") is not None + + +# Lazy imports to avoid import errors when strands is not installed +def __getattr__(name: str) -> Any: + """Lazy import of integration components.""" + if name == "HeadroomHookProvider": + from .hooks import HeadroomHookProvider + + return HeadroomHookProvider + elif name == "HeadroomStrandsModel": + from .model import HeadroomStrandsModel + + return HeadroomStrandsModel + elif name == "OptimizationMetrics": + from .model import OptimizationMetrics + + return OptimizationMetrics + elif name == "optimize_messages": + from .model import optimize_messages + + return optimize_messages + elif name == "get_headroom_provider": + from .providers import get_headroom_provider + + return get_headroom_provider + elif name == "get_model_name_from_strands": + from .providers import get_model_name_from_strands + + return get_model_name_from_strands + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = [ + # Availability check + "strands_available", + # Hook provider + "HeadroomHookProvider", + # Model wrapper + "HeadroomStrandsModel", + "OptimizationMetrics", + "optimize_messages", + # Provider detection + "get_headroom_provider", + "get_model_name_from_strands", +] diff --git a/headroom/integrations/strands/hooks.py b/headroom/integrations/strands/hooks.py new file mode 100644 index 000000000..29dc32a03 --- /dev/null +++ b/headroom/integrations/strands/hooks.py @@ -0,0 +1,540 @@ +"""Strands SDK hook provider for Headroom tool output compression. + +This module provides HeadroomHookProvider, which implements Strands' HookProvider +interface to intercept tool outputs and compress them using Headroom's SmartCrusher. + +Example: + from strands import Agent + from headroom.integrations.strands import HeadroomHookProvider + + # Create the hook provider + hook_provider = HeadroomHookProvider( + compress_tool_outputs=True, + min_tokens_to_compress=100, + ) + + # Use with Strands agent + agent = Agent(hooks=[hook_provider]) + response = agent("Search for documents about AI") + + # Check compression metrics + print(f"Tokens saved: {hook_provider.total_tokens_saved}") +""" + +from __future__ import annotations + +import json +import logging +import threading +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any +from uuid import uuid4 + +# Strands imports - these are optional dependencies +try: + from strands.hooks import HookProvider, HookRegistry + from strands.hooks.events import AfterToolCallEvent, BeforeToolCallEvent + from strands.types.tools import ToolResult + + STRANDS_AVAILABLE = True +except ImportError: + STRANDS_AVAILABLE = False + # Type stubs for when strands is not installed + HookProvider = object # type: ignore[misc,assignment] + HookRegistry = object # type: ignore[misc,assignment] + AfterToolCallEvent = object # type: ignore[misc,assignment] + BeforeToolCallEvent = object # type: ignore[misc,assignment] + ToolResult = dict # type: ignore[misc,assignment] + +from headroom import HeadroomConfig +from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig + +logger = logging.getLogger(__name__) + + +def _check_strands_available() -> None: + """Raise ImportError if Strands is not installed.""" + if not STRANDS_AVAILABLE: + raise ImportError( + "Strands SDK is required for this integration. Install with: pip install strands-agents" + ) + + +def strands_available() -> bool: + """Check if Strands SDK is installed. + + Returns: + True if strands-agents package is available. + """ + return STRANDS_AVAILABLE + + +@dataclass +class CompressionMetrics: + """Metrics from a single tool output compression.""" + + request_id: str + timestamp: datetime + tool_name: str + tool_use_id: str + tokens_before: int + tokens_after: int + tokens_saved: int + savings_percent: float + was_compressed: bool + skip_reason: str | None = None + + +@dataclass +class HeadroomHookProvider(HookProvider): # type: ignore[misc] + """Strands HookProvider that compresses tool outputs using Headroom. + + This hook provider intercepts tool call results via AfterToolCallEvent + and applies Headroom's SmartCrusher to compress large outputs, reducing + token usage while preserving important information. + + The compression is intelligent and preserves: + - Error items (containing error indicators) + - Anomalous values (statistical outliers) + - Items matching the user's query context + - First/last items for context + - Structural outliers (rare status values) + + Attributes: + compress_tool_outputs: Whether to compress tool outputs. + min_tokens_to_compress: Minimum token count before compression is applied. + config: Headroom configuration. + preserve_errors: If True, never compress results with error status. + total_tokens_saved: Running total of tokens saved across all compressions. + metrics_history: List of CompressionMetrics from recent compressions. + + Example: + from strands import Agent + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider(min_tokens_to_compress=50) + agent = Agent(hooks=[hook]) + + # After running agent tasks... + summary = hook.get_savings_summary() + print(f"Total saved: {summary['total_tokens_saved']} tokens") + """ + + compress_tool_outputs: bool = True + min_tokens_to_compress: int = 100 + config: HeadroomConfig | None = field(default=None) + preserve_errors: bool = True + + # Internal state (not part of dataclass comparison) + _crusher: SmartCrusher | None = field(default=None, repr=False, compare=False) + _metrics_history: list[CompressionMetrics] = field( + default_factory=list, repr=False, compare=False + ) + _total_tokens_saved: int = field(default=0, repr=False, compare=False) + _lock: threading.Lock = field(default_factory=threading.Lock, repr=False, compare=False) + _initialized: bool = field(default=False, repr=False, compare=False) + + def __post_init__(self) -> None: + """Initialize the hook provider after dataclass construction.""" + _check_strands_available() + + if self.config is None: + self.config = HeadroomConfig() + + self._initialized = True + logger.debug( + "HeadroomHookProvider initialized: compress=%s, min_tokens=%d, preserve_errors=%s", + self.compress_tool_outputs, + self.min_tokens_to_compress, + self.preserve_errors, + ) + + @property + def crusher(self) -> SmartCrusher: + """Lazily initialize SmartCrusher (thread-safe). + + Returns: + The SmartCrusher instance for compression. + """ + if self._crusher is None: + with self._lock: + # Double-check after acquiring lock + if self._crusher is None: + # Use config from HeadroomConfig if available + if self.config and self.config.smart_crusher: + crusher_config = SmartCrusherConfig( + min_tokens_to_crush=self.min_tokens_to_compress, + max_items_after_crush=self.config.smart_crusher.max_items_after_crush, + ) + else: + crusher_config = SmartCrusherConfig( + min_tokens_to_crush=self.min_tokens_to_compress + ) + self._crusher = SmartCrusher(config=crusher_config) + logger.debug( + "SmartCrusher initialized with min_tokens=%d", self.min_tokens_to_compress + ) + return self._crusher + + @property + def total_tokens_saved(self) -> int: + """Total tokens saved across all compressions. + + Returns: + Cumulative token savings. + """ + return self._total_tokens_saved + + @property + def metrics_history(self) -> list[CompressionMetrics]: + """History of compression metrics. + + Returns: + Copy of the metrics history list. + """ + return self._metrics_history.copy() + + def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None: + """Register hooks with the Strands HookRegistry. + + This method is called by Strands when the hook provider is added + to an Agent. It registers the compression handler for AfterToolCallEvent. + + Args: + registry: The Strands HookRegistry to register hooks with. + """ + if not self.compress_tool_outputs: + logger.debug("Tool output compression disabled, skipping hook registration") + return + + # Register the after-tool-call hook for compression + registry.add_callback(AfterToolCallEvent, self._compress_tool_result) + logger.info( + "HeadroomHookProvider registered: compressing tool outputs >= %d tokens", + self.min_tokens_to_compress, + ) + + def _estimate_tokens(self, text: str) -> int: + """Estimate token count for text. + + Uses a simple heuristic of ~4 characters per token, which is + reasonably accurate for English text and JSON content. + + Args: + text: The text to estimate tokens for. + + Returns: + Estimated token count. + """ + if not text: + return 0 + # ~4 characters per token is a reasonable estimate + return len(text) // 4 + + def _extract_text_content(self, result: ToolResult) -> str: + """Extract text content from a ToolResult. + + Handles both text and JSON content types in the result. + + Args: + result: The ToolResult to extract content from. + + Returns: + String representation of the content. + """ + content = result.get("content", []) + if not content: + return "" + + text_parts = [] + for item in content: + if isinstance(item, dict): + if "text" in item: + text_parts.append(str(item["text"])) + elif "json" in item: + try: + text_parts.append(json.dumps(item["json"], indent=None)) + except (TypeError, ValueError): + text_parts.append(str(item["json"])) + elif isinstance(item, str): + text_parts.append(item) + + return "\n".join(text_parts) + + def _update_result_content(self, result: ToolResult, compressed_text: str) -> None: + """Update the result content with compressed text. + + Modifies the result in place, preserving the original content structure + (text vs json) where possible. + + Args: + result: The ToolResult to update (modified in place). + compressed_text: The compressed content to set. + """ + content = result.get("content", []) + + if not content: + # No existing content, create text content + result["content"] = [{"text": compressed_text}] + return + + # Try to preserve original structure + first_item = content[0] if content else None + + if isinstance(first_item, dict): + if "json" in first_item: + # Try to parse compressed text back to JSON + try: + parsed = json.loads(compressed_text) + result["content"] = [{"json": parsed}] + except (json.JSONDecodeError, ValueError): + # Fall back to text if not valid JSON + result["content"] = [{"text": compressed_text}] + else: + # Text content + result["content"] = [{"text": compressed_text}] + else: + # Unknown structure, use text + result["content"] = [{"text": compressed_text}] + + def _compress_tool_result(self, event: AfterToolCallEvent) -> None: + """Compress tool result content if it exceeds the token threshold. + + This is the main hook handler that intercepts AfterToolCallEvent + and applies SmartCrusher compression to large tool outputs. + + Args: + event: The AfterToolCallEvent containing the tool result. + The result field is writable and modified in place. + """ + request_id = str(uuid4()) + result = event.result + tool_name = event.tool_use.get("name", "unknown") + tool_use_id = event.tool_use.get("toolUseId", "unknown") + + # Check if compression should be skipped + skip_reason = self._should_skip_compression(result) + if skip_reason: + self._record_metrics( + request_id=request_id, + tool_name=tool_name, + tool_use_id=tool_use_id, + tokens_before=0, + tokens_after=0, + was_compressed=False, + skip_reason=skip_reason, + ) + logger.debug( + "Skipping compression for tool %s (id=%s): %s", + tool_name, + tool_use_id, + skip_reason, + ) + return + + # Extract content and estimate tokens + original_text = self._extract_text_content(result) + tokens_before = self._estimate_tokens(original_text) + + # Check minimum token threshold + if tokens_before < self.min_tokens_to_compress: + self._record_metrics( + request_id=request_id, + tool_name=tool_name, + tool_use_id=tool_use_id, + tokens_before=tokens_before, + tokens_after=tokens_before, + was_compressed=False, + skip_reason=f"below_threshold:{tokens_before}<{self.min_tokens_to_compress}", + ) + logger.debug( + "Tool %s output below threshold (%d < %d tokens), skipping compression", + tool_name, + tokens_before, + self.min_tokens_to_compress, + ) + return + + # Apply compression + try: + crush_result = self.crusher.crush(content=original_text, query="") + compressed_text = crush_result.compressed + was_modified = crush_result.was_modified + except Exception as e: + # Compression failed, keep original + logger.warning( + "Compression failed for tool %s (id=%s): %s. Keeping original.", + tool_name, + tool_use_id, + str(e), + ) + self._record_metrics( + request_id=request_id, + tool_name=tool_name, + tool_use_id=tool_use_id, + tokens_before=tokens_before, + tokens_after=tokens_before, + was_compressed=False, + skip_reason=f"compression_error:{type(e).__name__}", + ) + return + + tokens_after = self._estimate_tokens(compressed_text) + + # Only update if compression actually reduced tokens + if was_modified and tokens_after < tokens_before: + self._update_result_content(result, compressed_text) + tokens_saved = tokens_before - tokens_after + + self._record_metrics( + request_id=request_id, + tool_name=tool_name, + tool_use_id=tool_use_id, + tokens_before=tokens_before, + tokens_after=tokens_after, + was_compressed=True, + skip_reason=None, + ) + + logger.info( + "Compressed tool %s output: %d -> %d tokens (%.1f%% saved)", + tool_name, + tokens_before, + tokens_after, + (tokens_saved / tokens_before * 100) if tokens_before > 0 else 0, + ) + else: + # Compression didn't help + self._record_metrics( + request_id=request_id, + tool_name=tool_name, + tool_use_id=tool_use_id, + tokens_before=tokens_before, + tokens_after=tokens_before, + was_compressed=False, + skip_reason="no_reduction", + ) + logger.debug( + "Compression did not reduce tool %s output (%d tokens)", + tool_name, + tokens_before, + ) + + def _should_skip_compression(self, result: ToolResult) -> str | None: + """Check if compression should be skipped for this result. + + Args: + result: The tool result to check. + + Returns: + Skip reason string if should skip, None if should compress. + """ + # Skip if compression is disabled + if not self.compress_tool_outputs: + return "compression_disabled" + + # Skip error results if preserve_errors is True + if self.preserve_errors and result.get("status") == "error": + return "error_result_preserved" + + # Skip empty results + content = result.get("content", []) + if not content: + return "empty_content" + + return None + + def _record_metrics( + self, + request_id: str, + tool_name: str, + tool_use_id: str, + tokens_before: int, + tokens_after: int, + was_compressed: bool, + skip_reason: str | None, + ) -> None: + """Record compression metrics (thread-safe). + + Args: + request_id: Unique ID for this compression request. + tool_name: Name of the tool that was called. + tool_use_id: The toolUseId from the result. + tokens_before: Token count before compression. + tokens_after: Token count after compression. + was_compressed: Whether compression was actually applied. + skip_reason: Reason compression was skipped, if applicable. + """ + tokens_saved = max(0, tokens_before - tokens_after) + savings_percent = (tokens_saved / tokens_before * 100) if tokens_before > 0 else 0.0 + + metrics = CompressionMetrics( + request_id=request_id, + timestamp=datetime.now(timezone.utc), + tool_name=tool_name, + tool_use_id=tool_use_id, + tokens_before=tokens_before, + tokens_after=tokens_after, + tokens_saved=tokens_saved, + savings_percent=savings_percent, + was_compressed=was_compressed, + skip_reason=skip_reason, + ) + + with self._lock: + self._metrics_history.append(metrics) + if was_compressed: + self._total_tokens_saved += tokens_saved + + # Keep only last 100 metrics to bound memory + if len(self._metrics_history) > 100: + self._metrics_history = self._metrics_history[-100:] + + def get_savings_summary(self) -> dict[str, Any]: + """Get summary of token savings across all compressions. + + Returns: + Dictionary with compression statistics including: + - total_requests: Number of tool outputs processed + - compressed_requests: Number actually compressed + - total_tokens_saved: Cumulative tokens saved + - average_savings_percent: Mean compression ratio + - total_tokens_before: Sum of all input tokens + - total_tokens_after: Sum of all output tokens + """ + if not self._metrics_history: + return { + "total_requests": 0, + "compressed_requests": 0, + "total_tokens_saved": 0, + "average_savings_percent": 0.0, + "total_tokens_before": 0, + "total_tokens_after": 0, + } + + compressed_metrics = [m for m in self._metrics_history if m.was_compressed] + + return { + "total_requests": len(self._metrics_history), + "compressed_requests": len(compressed_metrics), + "total_tokens_saved": self._total_tokens_saved, + "average_savings_percent": ( + sum(m.savings_percent for m in compressed_metrics) / len(compressed_metrics) + if compressed_metrics + else 0.0 + ), + "total_tokens_before": sum(m.tokens_before for m in self._metrics_history), + "total_tokens_after": sum(m.tokens_after for m in self._metrics_history), + } + + def reset(self) -> None: + """Reset all tracked metrics (thread-safe). + + Clears the metrics history and resets the total tokens saved counter. + Useful for starting fresh measurements or between test runs. + """ + with self._lock: + self._metrics_history = [] + self._total_tokens_saved = 0 + logger.debug("HeadroomHookProvider metrics reset") diff --git a/headroom/integrations/strands/model.py b/headroom/integrations/strands/model.py new file mode 100644 index 000000000..abc62e08b --- /dev/null +++ b/headroom/integrations/strands/model.py @@ -0,0 +1,625 @@ +"""Strands SDK model wrapper for Headroom optimization. + +This module provides HeadroomStrandsModel, which wraps any Strands model +to apply Headroom context optimization before API calls. +""" + +from __future__ import annotations + +import asyncio +import logging +import threading +from collections.abc import AsyncGenerator, AsyncIterable +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any, TypeVar +from uuid import uuid4 + +# Strands imports - these are optional dependencies +try: + from strands.models import Model + from strands.types.content import Message, Messages, SystemContentBlock + from strands.types.streaming import StreamEvent + from strands.types.tools import ToolChoice, ToolSpec + + STRANDS_AVAILABLE = True +except ImportError: + STRANDS_AVAILABLE = False + Model = object # type: ignore[misc,assignment] + Message = dict # type: ignore[misc,assignment] + Messages = list # type: ignore[misc,assignment] + StreamEvent = dict # type: ignore[misc,assignment] + ToolChoice = dict # type: ignore[misc,assignment] + ToolSpec = dict # type: ignore[misc,assignment] + SystemContentBlock = dict # type: ignore[misc,assignment] + +T = TypeVar("T") + +from headroom import HeadroomConfig # noqa: E402 +from headroom.providers import OpenAIProvider # noqa: E402 +from headroom.transforms import TransformPipeline # noqa: E402 + +from .providers import get_headroom_provider, get_model_name_from_strands # noqa: E402 + +logger = logging.getLogger(__name__) + + +def _check_strands_available() -> None: + """Raise ImportError if Strands SDK is not installed.""" + if not STRANDS_AVAILABLE: + raise ImportError( + "Strands SDK is required for this integration. Install with: pip install strands-agents" + ) + + +def strands_available() -> bool: + """Check if Strands SDK is installed.""" + return STRANDS_AVAILABLE + + +@dataclass +class OptimizationMetrics: + """Metrics from a single optimization pass.""" + + request_id: str + timestamp: datetime + tokens_before: int + tokens_after: int + tokens_saved: int + savings_percent: float + transforms_applied: list[str] + model: str + + +class HeadroomStrandsModel(Model): # type: ignore[misc] + """Strands model wrapper that applies Headroom optimizations. + + Wraps any Strands Model and automatically optimizes the context + before each API call. Works with any Strands-compatible model provider. + + Example: + from strands import Agent + from strands.models.bedrock import BedrockModel + from headroom.integrations.strands import HeadroomStrandsModel + + # Basic usage + model = BedrockModel(model_id="us.anthropic.claude-sonnet-4-20250514-v1:0") + optimized = HeadroomStrandsModel(wrapped_model=model) + + # Use with agent + agent = Agent(model=optimized) + response = agent("Hello!") + + # Access metrics + print(f"Saved {optimized.total_tokens_saved} tokens") + + # With custom config + from headroom import HeadroomConfig + config = HeadroomConfig() + optimized = HeadroomStrandsModel(wrapped_model=model, config=config) + + Attributes: + wrapped_model: The underlying Strands model + total_tokens_saved: Running total of tokens saved + metrics_history: List of OptimizationMetrics from recent calls + """ + + def __init__( + self, + wrapped_model: Any, + config: HeadroomConfig | None = None, + auto_detect_provider: bool = True, + ) -> None: + """Initialize HeadroomStrandsModel. + + Args: + wrapped_model: The Strands model to wrap (e.g., BedrockModel, OpenAIModel) + config: Optional HeadroomConfig for optimization settings + auto_detect_provider: Whether to auto-detect the Headroom provider + based on the wrapped model type. Default True. + """ + _check_strands_available() + + if wrapped_model is None: + raise ValueError("wrapped_model cannot be None") + + self.wrapped_model = wrapped_model + self.headroom_config = config or HeadroomConfig() + self.auto_detect_provider = auto_detect_provider + + # Internal state + self._metrics_history: list[OptimizationMetrics] = [] + self._total_tokens_saved: int = 0 + self._pipeline: TransformPipeline | None = None + self._headroom_provider: Any = None + self._lock = threading.Lock() + + @property + def config(self) -> Any: + """Forward config access to wrapped model (required by Strands Agent).""" + return self.wrapped_model.config + + @property + def pipeline(self) -> TransformPipeline: + """Lazily initialize TransformPipeline (thread-safe).""" + if self._pipeline is None: + with self._lock: + # Double-check after acquiring lock + if self._pipeline is None: + if self.auto_detect_provider: + self._headroom_provider = get_headroom_provider(self.wrapped_model) + logger.debug( + f"Auto-detected provider: {self._headroom_provider.__class__.__name__}" + ) + else: + self._headroom_provider = OpenAIProvider() + self._pipeline = TransformPipeline( + config=self.headroom_config, + provider=self._headroom_provider, + ) + return self._pipeline + + @property + def total_tokens_saved(self) -> int: + """Total tokens saved across all calls.""" + return self._total_tokens_saved + + @property + def metrics_history(self) -> list[OptimizationMetrics]: + """History of optimization metrics.""" + return self._metrics_history.copy() + + def _convert_messages_to_openai(self, messages: list[Any]) -> list[dict[str, Any]]: + """Convert Strands messages to OpenAI format for Headroom. + + Strands uses dict-based messages similar to OpenAI format: + - {"role": "user", "content": "..."} + - {"role": "assistant", "content": "...", "tool_calls": [...]} + - {"role": "tool", "content": "...", "tool_call_id": "..."} + + Args: + messages: List of Strands messages (typically dicts or Message objects) + + Returns: + List of messages in OpenAI dict format + """ + result = [] + for msg in messages: + # Handle dict format (most common in Strands) + if isinstance(msg, dict): + entry: dict[str, Any] = { + "role": msg.get("role", "user"), + } + + # Handle content + content = msg.get("content") + if content is None: + entry["content"] = "" + elif isinstance(content, list): + # Content blocks - preserve structure + entry["content"] = content + else: + entry["content"] = content + + # Handle tool calls + if "tool_calls" in msg and msg["tool_calls"]: + entry["tool_calls"] = msg["tool_calls"] + + # Handle tool call ID for tool responses + if "tool_call_id" in msg and msg["tool_call_id"]: + entry["tool_call_id"] = msg["tool_call_id"] + + # Handle name field (for tool messages) + if "name" in msg and msg["name"]: + entry["name"] = msg["name"] + + result.append(entry) + + # Handle Strands Message objects (if they have role/content attrs) + elif hasattr(msg, "role") and hasattr(msg, "content"): + entry = { + "role": msg.role, + } + + content = msg.content + if content is None: + entry["content"] = "" + elif isinstance(content, list): + entry["content"] = content + else: + entry["content"] = content + + if hasattr(msg, "tool_calls") and msg.tool_calls: + entry["tool_calls"] = msg.tool_calls + if hasattr(msg, "tool_call_id") and msg.tool_call_id: + entry["tool_call_id"] = msg.tool_call_id + if hasattr(msg, "name") and msg.name: + entry["name"] = msg.name + + result.append(entry) + + else: + # Fallback: convert to string + content = str(msg) if msg is not None else "" + result.append({"role": "user", "content": content}) + + return result + + def _convert_messages_from_openai( + self, messages: list[dict[str, Any]], original_messages: list[Any] + ) -> list[dict[str, Any]]: + """Convert OpenAI format messages back to Strands format. + + Since Strands uses dict-based messages similar to OpenAI, + this is largely a passthrough, but ensures proper structure. + + Args: + messages: The optimized messages in OpenAI dict format + original_messages: The original Strands messages (for reference) + + Returns: + List of messages in Strands dict format + """ + result = [] + for msg in messages: + entry: dict[str, Any] = { + "role": msg.get("role", "user"), + } + + # Handle content + content = msg.get("content") + if content is not None: + entry["content"] = content + + # Preserve tool-related fields + if "tool_calls" in msg and msg["tool_calls"]: + entry["tool_calls"] = msg["tool_calls"] + if "tool_call_id" in msg and msg["tool_call_id"]: + entry["tool_call_id"] = msg["tool_call_id"] + if "name" in msg and msg["name"]: + entry["name"] = msg["name"] + + result.append(entry) + + return result + + def _optimize_messages( + self, messages: list[Any] + ) -> tuple[list[dict[str, Any]], OptimizationMetrics]: + """Apply Headroom optimization to messages. + + Thread-safe with fallback on pipeline errors. + + Args: + messages: List of Strands messages to optimize + + Returns: + Tuple of (optimized_messages, metrics) + """ + request_id = str(uuid4()) + + # Convert to OpenAI format + openai_messages = self._convert_messages_to_openai(messages) + + # Handle empty messages gracefully + if not openai_messages: + metrics = OptimizationMetrics( + request_id=request_id, + timestamp=datetime.now(timezone.utc), + tokens_before=0, + tokens_after=0, + tokens_saved=0, + savings_percent=0, + transforms_applied=[], + model=get_model_name_from_strands(self.wrapped_model), + ) + return [], metrics + + # Get model name from wrapped model + model = get_model_name_from_strands(self.wrapped_model) + + # Ensure pipeline is initialized + _ = self.pipeline + + # Get model context limit + model_limit = ( + self._headroom_provider.get_context_limit(model) if self._headroom_provider else 128000 + ) + + try: + # Apply Headroom transforms via pipeline + result = self.pipeline.apply( + messages=openai_messages, + model=model, + model_limit=model_limit, + ) + optimized = result.messages + tokens_before = result.tokens_before + tokens_after = result.tokens_after + transforms_applied = result.transforms_applied + except ( + ValueError, + TypeError, + AttributeError, + RuntimeError, + KeyError, + IndexError, + ImportError, + OSError, + ) as e: + # Fallback to original messages on pipeline error + logger.warning( + f"Headroom optimization failed, using original messages: {type(e).__name__}: {e}" + ) + optimized = openai_messages + # Estimate token count (rough approximation: ~4 chars/token) + tokens_before = sum(len(str(m.get("content", ""))) // 4 for m in openai_messages) + tokens_after = tokens_before + transforms_applied = ["fallback:error"] + + # Create metrics + tokens_saved = max(0, tokens_before - tokens_after) + metrics = OptimizationMetrics( + request_id=request_id, + timestamp=datetime.now(timezone.utc), + tokens_before=tokens_before, + tokens_after=tokens_after, + tokens_saved=tokens_saved, + savings_percent=(tokens_saved / tokens_before * 100 if tokens_before > 0 else 0), + transforms_applied=transforms_applied, + model=model, + ) + + # Track metrics (thread-safe) + with self._lock: + self._metrics_history.append(metrics) + self._total_tokens_saved += metrics.tokens_saved + + # Keep only last 100 metrics + if len(self._metrics_history) > 100: + self._metrics_history = self._metrics_history[-100:] + + # Convert back to Strands format + optimized_messages = self._convert_messages_from_openai(optimized, messages) + + return optimized_messages, metrics + + async def stream( + self, + messages: Messages, + tool_specs: list[ToolSpec] | None = None, + system_prompt: str | None = None, + *, + tool_choice: ToolChoice | None = None, + system_prompt_content: list[SystemContentBlock] | None = None, + invocation_state: dict[str, Any] | None = None, + **kwargs: Any, + ) -> AsyncIterable[StreamEvent]: + """Stream response with Headroom optimization. + + This is the main method required by Strands Model interface. + Optimizes messages before delegating to the wrapped model's stream method. + + Args: + messages: List of messages to send to the model + tool_specs: Optional list of tool specifications + system_prompt: Optional system prompt string + tool_choice: Optional tool choice configuration + system_prompt_content: Optional list of system content blocks + invocation_state: Optional invocation state dictionary + **kwargs: Additional arguments passed to the wrapped model + + Yields: + Streaming events from the wrapped model + """ + # Run optimization in executor (CPU-bound) + loop = asyncio.get_running_loop() + optimized_messages, metrics = await loop.run_in_executor( + None, self._optimize_messages, messages + ) + + logger.info( + f"Headroom optimized (stream): {metrics.tokens_before} -> " + f"{metrics.tokens_after} tokens ({metrics.savings_percent:.1f}% saved)" + ) + + # Delegate to wrapped model's stream method with all parameters + async for event in self.wrapped_model.stream( + optimized_messages, + tool_specs=tool_specs, + system_prompt=system_prompt, + tool_choice=tool_choice, + system_prompt_content=system_prompt_content, + invocation_state=invocation_state, + **kwargs, + ): + yield event + + def get_config(self) -> Any: + """Get the configuration of the wrapped model. + + Returns: + The model configuration from the wrapped model. + """ + return self.wrapped_model.get_config() + + def update_config(self, **model_config: Any) -> None: + """Update the configuration of the wrapped model. + + Args: + **model_config: Configuration options to update on the wrapped model. + """ + self.wrapped_model.update_config(**model_config) + + async def structured_output( + self, + output_model: type[T], + prompt: Messages, + system_prompt: str | None = None, + **kwargs: Any, + ) -> AsyncGenerator[dict[str, T | Any], None]: + """Generate structured output with Headroom optimization. + + Optimizes the prompt messages before delegating to the wrapped model's + structured_output method. + + Args: + output_model: The type/schema for the structured output + prompt: List of prompt messages + system_prompt: Optional system prompt + **kwargs: Additional arguments passed to the wrapped model + + Yields: + Structured output events from the wrapped model + """ + # Run optimization in executor (CPU-bound) + loop = asyncio.get_running_loop() + optimized_prompt, metrics = await loop.run_in_executor( + None, self._optimize_messages, prompt + ) + + logger.info( + f"Headroom optimized (structured_output): {metrics.tokens_before} -> " + f"{metrics.tokens_after} tokens ({metrics.savings_percent:.1f}% saved)" + ) + + # Delegate to wrapped model + async for event in self.wrapped_model.structured_output( + output_model, optimized_prompt, system_prompt=system_prompt, **kwargs + ): + yield event + + def get_savings_summary(self) -> dict[str, Any]: + """Get summary of token savings.""" + if not self._metrics_history: + return { + "total_requests": 0, + "total_tokens_saved": 0, + "average_savings_percent": 0, + "total_tokens_before": 0, + "total_tokens_after": 0, + } + + return { + "total_requests": len(self._metrics_history), + "total_tokens_saved": self._total_tokens_saved, + "average_savings_percent": sum(m.savings_percent for m in self._metrics_history) + / len(self._metrics_history), + "total_tokens_before": sum(m.tokens_before for m in self._metrics_history), + "total_tokens_after": sum(m.tokens_after for m in self._metrics_history), + } + + def reset(self) -> None: + """Reset all tracked metrics (thread-safe). + + Clears the metrics history and resets the total tokens saved counter. + Useful for starting fresh measurements or between test runs. + """ + with self._lock: + self._metrics_history = [] + self._total_tokens_saved = 0 + + # ========================================================================= + # Forward attribute access to wrapped model for compatibility + # ========================================================================= + + def __getattr__(self, name: str) -> Any: + """Forward attribute access to wrapped model.""" + # Avoid infinite recursion for our own attributes + if name in ( + "wrapped_model", + "config", + "auto_detect_provider", + "_metrics_history", + "_total_tokens_saved", + "_pipeline", + "_headroom_provider", + "_lock", + "pipeline", + "total_tokens_saved", + "metrics_history", + ): + raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'") + return getattr(self.wrapped_model, name) + + +def optimize_messages( + messages: list[Any], + config: HeadroomConfig | None = None, + model: str = "gpt-4o", +) -> tuple[list[dict[str, Any]], dict[str, Any]]: + """Standalone function to optimize Strands messages. + + Use this for manual optimization when you need fine-grained control. + + Args: + messages: List of Strands messages (dicts) + config: HeadroomConfig for optimization settings + model: Model name for token estimation + + Returns: + Tuple of (optimized_messages, metrics_dict) + + Example: + from headroom.integrations.strands import optimize_messages + + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "What is 2+2?"}, + ] + + optimized, metrics = optimize_messages(messages) + print(f"Saved {metrics['tokens_saved']} tokens") + """ + _check_strands_available() + + config = config or HeadroomConfig() + provider = OpenAIProvider() + pipeline = TransformPipeline(config=config, provider=provider) + + # Convert to OpenAI format (Strands uses similar format) + openai_messages = [] + for msg in messages: + if isinstance(msg, dict): + entry: dict[str, Any] = { + "role": msg.get("role", "user"), + "content": msg.get("content", ""), + } + if "tool_calls" in msg and msg["tool_calls"]: + entry["tool_calls"] = msg["tool_calls"] + if "tool_call_id" in msg and msg["tool_call_id"]: + entry["tool_call_id"] = msg["tool_call_id"] + openai_messages.append(entry) + elif hasattr(msg, "role") and hasattr(msg, "content"): + entry = {"role": msg.role, "content": msg.content or ""} + if hasattr(msg, "tool_calls") and msg.tool_calls: + entry["tool_calls"] = msg.tool_calls + if hasattr(msg, "tool_call_id") and msg.tool_call_id: + entry["tool_call_id"] = msg.tool_call_id + openai_messages.append(entry) + else: + openai_messages.append({"role": "user", "content": str(msg)}) + + # Get model context limit + model_limit = provider.get_context_limit(model) + + # Apply transforms + result = pipeline.apply( + messages=openai_messages, + model=model, + model_limit=model_limit, + ) + + metrics = { + "tokens_before": result.tokens_before, + "tokens_after": result.tokens_after, + "tokens_saved": result.tokens_before - result.tokens_after, + "savings_percent": ( + (result.tokens_before - result.tokens_after) / result.tokens_before * 100 + if result.tokens_before > 0 + else 0 + ), + "transforms_applied": result.transforms_applied, + } + + return result.messages, metrics diff --git a/headroom/integrations/strands/providers.py b/headroom/integrations/strands/providers.py new file mode 100644 index 000000000..f93609f1c --- /dev/null +++ b/headroom/integrations/strands/providers.py @@ -0,0 +1,166 @@ +"""Provider detection for Strands models. + +Automatically detects the correct Headroom provider based on the Strands model type. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from headroom.providers import ( + AnthropicProvider, + GoogleProvider, + OpenAIProvider, +) +from headroom.providers.base import Provider + +logger = logging.getLogger(__name__) + +# Mapping from Strands model class names to Headroom providers +_STRANDS_MODEL_PROVIDERS: dict[str, type[Provider]] = { + # Bedrock models (primarily Claude via Bedrock) + "BedrockModel": AnthropicProvider, + # Anthropic models (direct API) + "AnthropicModel": AnthropicProvider, + # OpenAI models + "OpenAIModel": OpenAIProvider, + # LiteLLM (uses OpenAI-compatible interface) + "LiteLLMModel": OpenAIProvider, + # Ollama (uses OpenAI-compatible interface) + "OllamaModel": OpenAIProvider, + # Google Gemini models + "GeminiModel": GoogleProvider, + # Writer models (uses OpenAI-compatible interface) + "WriterModel": OpenAIProvider, +} + + +def get_headroom_provider(model: Any) -> Provider: + """Get the appropriate Headroom provider for a Strands model. + + Detection strategy: + 1. Check model class name against known Strands model types + 2. Check for provider hints in model attributes + 3. Fall back to OpenAI provider (most compatible) + + Args: + model: A Strands model instance (BedrockModel, AnthropicModel, etc.) + + Returns: + Appropriate Headroom Provider instance. + + Example: + from strands.models import BedrockModel + from headroom.integrations.strands.providers import get_headroom_provider + + model = BedrockModel(model_id="anthropic.claude-3-5-sonnet-20241022-v2:0") + provider = get_headroom_provider(model) # Returns AnthropicProvider + """ + # Strategy 1: Class name matching + class_name = model.__class__.__name__ + if class_name in _STRANDS_MODEL_PROVIDERS: + provider_class = _STRANDS_MODEL_PROVIDERS[class_name] + logger.debug(f"Detected provider {provider_class.__name__} from class {class_name}") + return provider_class() + + # Strategy 2: Check module path + module_path = model.__class__.__module__ + if "anthropic" in module_path.lower(): + logger.debug(f"Detected AnthropicProvider from module {module_path}") + return AnthropicProvider() + elif "bedrock" in module_path.lower(): + logger.debug(f"Detected AnthropicProvider from module {module_path}") + return AnthropicProvider() + elif "google" in module_path.lower() or "gemini" in module_path.lower(): + logger.debug(f"Detected GoogleProvider from module {module_path}") + return GoogleProvider() + elif "openai" in module_path.lower() or "litellm" in module_path.lower(): + logger.debug(f"Detected OpenAIProvider from module {module_path}") + return OpenAIProvider() + + # Strategy 3: Check model ID/name for hints + model_id = _extract_model_id(model) + if model_id: + model_id_lower = model_id.lower() + if "claude" in model_id_lower or "anthropic" in model_id_lower: + logger.debug(f"Detected AnthropicProvider from model ID {model_id}") + return AnthropicProvider() + elif "gemini" in model_id_lower: + logger.debug(f"Detected GoogleProvider from model ID {model_id}") + return GoogleProvider() + elif "gpt" in model_id_lower or "o1" in model_id_lower or "o3" in model_id_lower: + logger.debug(f"Detected OpenAIProvider from model ID {model_id}") + return OpenAIProvider() + + # Strategy 4: Default fallback + logger.warning( + f"Unknown Strands model class '{class_name}', defaulting to OpenAIProvider. " + "Token counting may be inaccurate." + ) + return OpenAIProvider() + + +def _extract_model_id(model: Any) -> str: + """Extract model ID from a Strands model using various attribute names. + + Args: + model: A Strands model instance + + Returns: + Model ID string or empty string if not found + """ + # Try common attribute names used by Strands models + for attr in ["model_id", "model", "model_name", "id"]: + value = getattr(model, attr, None) + if value and isinstance(value, str): + return str(value) + + # Try to get from config if available (config can be dict or object) + config = getattr(model, "config", None) + if config: + for attr in ["model_id", "model", "model_name"]: + # Handle dict-style config (Strands uses this) + if isinstance(config, dict): + value = config.get(attr) + else: + value = getattr(config, attr, None) + if value and isinstance(value, str): + return str(value) + + # Try get_config() method (Strands Model interface) + if hasattr(model, "get_config"): + try: + config_dict = model.get_config() + if isinstance(config_dict, dict): + for attr in ["model_id", "model", "model_name"]: + value = config_dict.get(attr) + if value and isinstance(value, str): + return str(value) + except Exception: + pass + + return "" + + +def get_model_name_from_strands(model: Any) -> str: + """Extract the model name/ID from a Strands model. + + Args: + model: A Strands model instance + + Returns: + Model name string (e.g., "anthropic.claude-3-5-sonnet-20241022-v2:0") + """ + model_id = _extract_model_id(model) + if model_id: + return str(model_id) + + # Fallback with warning + class_name = model.__class__.__name__ + logger.warning( + f"Could not extract model name from {class_name} (no 'model_id', 'model', " + f"'model_name', or 'id' attribute). Defaulting to 'gpt-4o'. " + "Token counting may be inaccurate." + ) + return "gpt-4o" diff --git a/pyproject.toml b/pyproject.toml index daab7a1b4..c894e308b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -90,6 +90,10 @@ code = [ agno = [ "agno>=1.0.0", ] +# AWS Strands Agents SDK integration +strands = [ + "strands-agents>=0.1.0", +] # Voice filler detection (training and inference) voice = [ "onnxruntime>=1.16.0", # Fast CPU inference diff --git a/tests/integrations/test_strands/__init__.py b/tests/integrations/test_strands/__init__.py new file mode 100644 index 000000000..f87b7b7bb --- /dev/null +++ b/tests/integrations/test_strands/__init__.py @@ -0,0 +1 @@ +"""Tests for Strands Agents SDK integration with Headroom.""" diff --git a/tests/integrations/test_strands/test_hooks.py b/tests/integrations/test_strands/test_hooks.py new file mode 100644 index 000000000..f88841e0d --- /dev/null +++ b/tests/integrations/test_strands/test_hooks.py @@ -0,0 +1,545 @@ +"""Real-world integration tests for Strands HeadroomHookProvider. + +These tests use actual AWS Bedrock API calls with real credentials. +NO MOCKS - all tests hit the real Bedrock API. + +Skip in CI if AWS credentials are not available. +""" + +from __future__ import annotations + +import json +import os + +import pytest + +# Check for AWS credentials availability +SKIP_BEDROCK = not ( + os.environ.get("AWS_ACCESS_KEY_ID") + or os.environ.get("AWS_PROFILE") + or os.path.exists(os.path.expanduser("~/.aws/credentials")) +) + +# Check if strands-agents is installed +try: + from strands import Agent, tool + from strands.models import BedrockModel + + STRANDS_AVAILABLE = True +except ImportError: + STRANDS_AVAILABLE = False + + # Provide a no-op decorator when strands is not installed + def tool(fn): + return fn + + Agent = None # type: ignore + BedrockModel = None # type: ignore + +# Skip all tests if dependencies not available +pytestmark = [ + pytest.mark.skipif(SKIP_BEDROCK, reason="AWS credentials not available"), + pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed"), +] + + +# ============================================================================ +# Test Tools - Generate realistic verbose data for compression testing +# These are defined with @tool decorator for use when strands is installed. +# When strands is not installed, the no-op decorator ensures import succeeds. +# ============================================================================ + + +@tool +def search_logs(query: str, limit: int = 100) -> str: + """Search application logs. Returns JSON array of log entries. + + Args: + query: Search query to find in logs + limit: Maximum number of log entries to return + + Returns: + JSON array of log entry objects + """ + # Generate realistic verbose log data that should be compressed + logs = [ + { + "timestamp": f"2024-01-{(i % 28) + 1:02d}T{10 + (i % 12):02d}:00:00Z", + "level": ["INFO", "DEBUG", "WARN", "ERROR"][i % 4], + "service": ["api-gateway", "auth-service", "data-processor", "cache-service"][i % 4], + "message": f"Request processed successfully - latency={50 + i}ms, query={query}", + "request_id": f"req-{i:06d}-{hash(query) % 10000:04d}", + "status_code": [200, 201, 400, 500][i % 4], + "user_agent": "Mozilla/5.0 (compatible; TestBot/1.0)", + "ip_address": f"192.168.{i % 256}.{(i * 7) % 256}", + "trace_id": f"trace-{i:08x}", + "span_id": f"span-{i:04x}", + "duration_ms": 50 + (i * 3) % 200, + "memory_mb": 128 + (i * 5) % 512, + "cpu_percent": 10 + (i * 2) % 80, + } + for i in range(limit) + ] + return json.dumps(logs, indent=2) + + +@tool +def get_small_status() -> str: + """Get a small status response that should NOT be compressed. + + Returns: + Small JSON status object + """ + return json.dumps({"status": "healthy", "uptime_seconds": 12345, "version": "1.2.3"}) + + +@tool +def get_error_data() -> str: + """Get error information. Error results should NOT be compressed. + + Returns: + Error information (but not as a tool error) + """ + return json.dumps( + { + "errors": [ + {"code": "E001", "message": "Connection timeout"}, + {"code": "E002", "message": "Authentication failed"}, + ], + "timestamp": "2024-01-15T10:00:00Z", + } + ) + + +@tool +def fetch_user_data(user_id: str) -> str: + """Fetch detailed user data. Returns large JSON payload. + + Args: + user_id: The user ID to fetch data for + + Returns: + Large JSON object with user details + """ + # Generate a large user profile that should trigger compression + activities = [ + { + "activity_id": f"act-{i:06d}", + "type": ["login", "purchase", "view", "share"][i % 4], + "timestamp": f"2024-01-{(i % 28) + 1:02d}T{10 + (i % 12):02d}:30:00Z", + "details": { + "ip": f"10.0.{i % 256}.{(i * 3) % 256}", + "device": ["desktop", "mobile", "tablet"][i % 3], + "browser": ["Chrome", "Firefox", "Safari"][i % 3], + "duration_seconds": 30 + i * 5, + "page_views": 1 + i % 10, + }, + "metadata": { + "session_id": f"sess-{i:08x}", + "referrer": f"https://example.com/page/{i}", + "utm_source": ["google", "facebook", "twitter", "email"][i % 4], + }, + } + for i in range(50) + ] + + return json.dumps( + { + "user_id": user_id, + "profile": { + "name": "Test User", + "email": f"{user_id}@example.com", + "created_at": "2023-01-01T00:00:00Z", + }, + "activities": activities, + }, + indent=2, + ) + + +@tool +def simple_calculator(a: int, b: int, operation: str) -> str: + """Simple calculator for basic operations. + + Args: + a: First number + b: Second number + operation: One of 'add', 'subtract', 'multiply', 'divide' + + Returns: + The result of the operation + """ + if operation == "add": + result = a + b + elif operation == "subtract": + result = a - b + elif operation == "multiply": + result = a * b + elif operation == "divide": + result = a / b if b != 0 else "undefined" + else: + result = "unknown operation" + + return json.dumps({"operation": operation, "a": a, "b": b, "result": result}) + + +# ============================================================================ +# Test Class +# ============================================================================ + + +@pytest.mark.skipif(SKIP_BEDROCK, reason="AWS credentials not available") +@pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed") +class TestHeadroomHookProviderReal: + """Real-world integration tests for HeadroomHookProvider with Bedrock.""" + + @pytest.fixture + def bedrock_model(self): + """Create a BedrockModel instance using Claude 3 Haiku (fast and cheap).""" + return BedrockModel( + model_id="anthropic.claude-3-haiku-20240307-v1:0", + region_name="us-west-2", + temperature=0.1, # Low temperature for consistent tests + ) + + @pytest.fixture + def hook_provider(self): + """Create a HeadroomHookProvider with test configuration.""" + from headroom.integrations.strands import HeadroomHookProvider + + return HeadroomHookProvider( + compress_tool_outputs=True, + min_tokens_to_compress=50, # Low threshold for testing + preserve_errors=True, + ) + + def test_hook_compresses_large_tool_output(self, bedrock_model, hook_provider): + """Test that large tool outputs are compressed by the hook. + + This test: + 1. Creates an agent with the search_logs tool + 2. Asks a question that triggers the tool + 3. Verifies the hook compressed the output and saved tokens + """ + # Create agent with hook provider + agent = Agent( + model=bedrock_model, + tools=[search_logs], + hooks=[hook_provider], + ) + + # Ask a question that will trigger the search_logs tool + result = agent( + "Search the logs for 'error' and tell me how many entries you found. " + "Use limit=100 to get plenty of results." + ) + + # Verify the agent got a response + assert result is not None + + # Check hook metrics + metrics = hook_provider.get_savings_summary() + + # The hook should have processed at least one tool call + assert metrics["total_requests"] >= 1, "Hook should have processed tool calls" + + # With 100 log entries, compression should have occurred + # and saved significant tokens + if metrics["compressed_requests"] > 0: + assert metrics["total_tokens_saved"] > 0, "Should have saved tokens" + assert metrics["total_tokens_before"] > metrics["total_tokens_after"] + + def test_hook_preserves_small_outputs(self, bedrock_model, hook_provider): + """Test that small tool outputs are NOT compressed. + + This test: + 1. Creates an agent with a tool returning small output + 2. Triggers the tool + 3. Verifies the hook did not modify the small output + """ + # Reset metrics from any previous tests + hook_provider.reset() + + agent = Agent( + model=bedrock_model, + tools=[get_small_status], + hooks=[hook_provider], + ) + + # Ask a question that will trigger the small status tool + result = agent("What is the current system status? Use the get_small_status tool.") + + assert result is not None + + # Check metrics - small outputs should not be compressed + metrics = hook_provider.get_savings_summary() + + # Tool was called but output was below threshold + if metrics["total_requests"] > 0: + # For small outputs, tokens_before == tokens_after (no compression) + for m in hook_provider.metrics_history: + if m.tool_name == "get_small_status" or "small" in str(m.skip_reason): + # Either not compressed or skip reason indicates below threshold + assert not m.was_compressed or m.skip_reason is not None, ( + "Small output should not be compressed" + ) + + def test_hook_preserves_errors(self, bedrock_model): + """Test that error results are NOT compressed when preserve_errors=True. + + This test: + 1. Creates a hook with preserve_errors=True + 2. Creates an agent with a tool that returns error data + 3. Verifies error results are preserved unchanged + """ + from headroom.integrations.strands import HeadroomHookProvider + + # Create hook with preserve_errors=True (default) + hook_with_preserve = HeadroomHookProvider( + compress_tool_outputs=True, + min_tokens_to_compress=10, # Very low threshold + preserve_errors=True, + ) + + agent = Agent( + model=bedrock_model, + tools=[get_error_data], + hooks=[hook_with_preserve], + ) + + # Get error data + result = agent("Get the error data using get_error_data tool and summarize it.") + + assert result is not None + + # Check that error-related results were handled appropriately + metrics = hook_with_preserve.get_savings_summary() + + # The get_error_data tool returns data about errors but doesn't itself error + # So it should be processed normally (this tests the flow works) + assert metrics["total_requests"] >= 0 # May or may not have been called + + def test_hook_metrics_tracking(self, bedrock_model, hook_provider): + """Test that metrics are tracked correctly across multiple tool calls. + + This test: + 1. Creates an agent with multiple tools + 2. Makes requests that trigger various tools + 3. Verifies metrics are accumulated correctly + """ + # Reset metrics + hook_provider.reset() + + agent = Agent( + model=bedrock_model, + tools=[search_logs, get_small_status, simple_calculator], + hooks=[hook_provider], + ) + + # First request - should trigger search_logs (large output) + agent("Search logs for 'test' with limit=50 and give me a count.") + + # Second request - should trigger calculator (small output) + agent("Calculate 15 + 27 using the calculator tool.") + + # Third request - should trigger status (small output) + agent("Get the system status using get_small_status.") + + # Check accumulated metrics + metrics = hook_provider.get_savings_summary() + + # Should have tracked multiple requests + assert metrics["total_requests"] >= 1, "Should have tracked tool requests" + + # total_tokens_before should be >= total_tokens_after + assert metrics["total_tokens_before"] >= metrics["total_tokens_after"] + + # History should contain records + history = hook_provider.metrics_history + assert len(history) >= 1, "Should have metrics history entries" + + # Each metric should have required fields + for m in history: + assert m.request_id is not None + assert m.timestamp is not None + assert m.tokens_before >= 0 + assert m.tokens_after >= 0 + + def test_multiple_tool_calls_in_single_request(self, bedrock_model, hook_provider): + """Test that multiple tool calls in a single agent request are all processed. + + This test: + 1. Asks a complex question requiring multiple tools + 2. Verifies each tool call is processed by the hook + """ + # Reset metrics + hook_provider.reset() + + agent = Agent( + model=bedrock_model, + tools=[search_logs, simple_calculator, fetch_user_data], + hooks=[hook_provider], + ) + + # Ask a complex question that might trigger multiple tools + result = agent( + "I need you to do three things: " + "1. Search logs for 'api' with limit=30. " + "2. Calculate 100 * 5 using the calculator. " + "3. Tell me the total number of results from step 1." + ) + + assert result is not None + + # Check that multiple tool calls were processed + metrics = hook_provider.get_savings_summary() + + # Should have processed at least the search_logs call + assert metrics["total_requests"] >= 1 + + # Verify metrics history + history = hook_provider.metrics_history + + # At minimum, should have processed search_logs (which has large output) + # The actual tools called depend on the model's interpretation + assert len(history) >= 1 + + # Check that we have tool names recorded + tool_names = [m.tool_name for m in history] + assert all(name is not None for name in tool_names) + + def test_hook_reset_clears_metrics(self, bedrock_model, hook_provider): + """Test that reset() clears all accumulated metrics. + + This test: + 1. Makes some requests to accumulate metrics + 2. Calls reset() + 3. Verifies all metrics are cleared + """ + agent = Agent( + model=bedrock_model, + tools=[search_logs], + hooks=[hook_provider], + ) + + # Make a request to accumulate metrics + agent("Search logs for 'test' with limit=20.") + + # Verify we have some metrics + assert hook_provider.total_tokens_saved >= 0 + + # Reset + hook_provider.reset() + + # Verify metrics are cleared + assert hook_provider.total_tokens_saved == 0 + assert len(hook_provider.metrics_history) == 0 + + metrics = hook_provider.get_savings_summary() + assert metrics["total_requests"] == 0 + assert metrics["total_tokens_saved"] == 0 + + def test_hook_with_compression_disabled(self, bedrock_model): + """Test that hook passes through without compression when disabled. + + This test: + 1. Creates a hook with compress_tool_outputs=False + 2. Verifies tool outputs are not modified + """ + from headroom.integrations.strands import HeadroomHookProvider + + # Create hook with compression disabled + disabled_hook = HeadroomHookProvider( + compress_tool_outputs=False, + min_tokens_to_compress=10, + ) + + agent = Agent( + model=bedrock_model, + tools=[search_logs], + hooks=[disabled_hook], + ) + + result = agent("Search logs for 'api' with limit=50.") + + assert result is not None + + # When compression is disabled, no requests should be tracked + # (the hook doesn't register callbacks when disabled) + metrics = disabled_hook.get_savings_summary() + assert metrics["compressed_requests"] == 0 + + def test_hook_concurrent_safety(self, bedrock_model, hook_provider): + """Test that hook is thread-safe for concurrent access. + + This test verifies that metrics tracking is thread-safe + by checking that accumulated values are consistent. + """ + import threading + + # Reset metrics + hook_provider.reset() + + agent = Agent( + model=bedrock_model, + tools=[simple_calculator], + hooks=[hook_provider], + ) + + results = [] + errors = [] + + def make_request(n: int): + try: + result = agent(f"Calculate {n} + {n} using simple_calculator.") + results.append(result) + except Exception as e: + errors.append(e) + + # Run a few sequential requests (concurrent Bedrock calls might be rate-limited) + threads = [] + for i in range(3): + t = threading.Thread(target=make_request, args=(i,)) + threads.append(t) + t.start() + # Small delay to avoid rate limiting + import time + + time.sleep(0.5) + + for t in threads: + t.join(timeout=60) # 60 second timeout per thread + + # Check we got results (some may have failed due to rate limits) + assert len(results) > 0 or len(errors) > 0 + + # Metrics should still be consistent + metrics = hook_provider.get_savings_summary() + assert metrics["total_tokens_before"] >= metrics["total_tokens_after"] + + def test_hook_handles_empty_tool_response(self, bedrock_model, hook_provider): + """Test that hook handles tools returning empty responses gracefully.""" + + @tool + def empty_response() -> str: + """Return an empty response.""" + return "" + + hook_provider.reset() + + agent = Agent( + model=bedrock_model, + tools=[empty_response], + hooks=[hook_provider], + ) + + # This might not trigger the tool if the model decides it's not needed + result = agent("Call the empty_response tool and tell me what you got.") + + assert result is not None + + # Should handle gracefully without errors + metrics = hook_provider.get_savings_summary() + # Just verify no exceptions and metrics are valid + assert metrics["total_tokens_before"] >= 0 + assert metrics["total_tokens_after"] >= 0 diff --git a/tests/integrations/test_strands/test_hooks_unit.py b/tests/integrations/test_strands/test_hooks_unit.py new file mode 100644 index 000000000..b059b93c5 --- /dev/null +++ b/tests/integrations/test_strands/test_hooks_unit.py @@ -0,0 +1,564 @@ +"""Unit tests for Strands HeadroomHookProvider. + +These tests use mocks and do NOT require AWS credentials or strands-agents. +They test the internal logic of HeadroomHookProvider in isolation. + +For real integration tests, see test_hooks.py. +""" + +from __future__ import annotations + +import json +import threading +from datetime import datetime, timezone +from unittest.mock import MagicMock + +import pytest + +# Check if strands-agents is installed for proper skip handling +try: + import strands # noqa: F401 + + STRANDS_AVAILABLE = True +except ImportError: + STRANDS_AVAILABLE = False + + +# Skip all tests if Strands not installed +pytestmark = pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed") + + +class TestHeadroomHookProviderInit: + """Tests for HeadroomHookProvider initialization.""" + + def test_init_with_defaults(self): + """Initialize with default settings.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + + assert hook.compress_tool_outputs is True + assert hook.min_tokens_to_compress == 100 + assert hook.preserve_errors is True + assert hook.total_tokens_saved == 0 + assert hook.metrics_history == [] + + def test_init_with_custom_config(self): + """Initialize with custom configuration.""" + from headroom import HeadroomConfig + from headroom.integrations.strands import HeadroomHookProvider + + config = HeadroomConfig() + config.smart_crusher.min_tokens_to_crush = 200 + config.smart_crusher.max_items_after_crush = 20 + + hook = HeadroomHookProvider( + compress_tool_outputs=False, + min_tokens_to_compress=500, + config=config, + preserve_errors=False, + ) + + assert hook.compress_tool_outputs is False + assert hook.min_tokens_to_compress == 500 + assert hook.config is config + assert hook.preserve_errors is False + + def test_init_creates_default_config_if_none(self): + """Initialize creates a default HeadroomConfig if none provided.""" + from headroom import HeadroomConfig + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + + assert hook.config is not None + assert isinstance(hook.config, HeadroomConfig) + + +class TestRegisterHooks: + """Tests for HeadroomHookProvider.register_hooks method.""" + + def test_register_hooks_adds_callback_to_registry(self): + """register_hooks adds AfterToolCallEvent callback to registry.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider(compress_tool_outputs=True) + mock_registry = MagicMock() + + hook.register_hooks(mock_registry) + + # Should have registered exactly one callback for AfterToolCallEvent + assert mock_registry.add_callback.call_count == 1 + + def test_register_hooks_skips_when_compression_disabled(self): + """register_hooks does not register callbacks when compression is disabled.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider(compress_tool_outputs=False) + mock_registry = MagicMock() + + hook.register_hooks(mock_registry) + + # Should not have registered any callbacks + assert mock_registry.add_callback.call_count == 0 + + +class TestCrusherLazyInit: + """Tests for SmartCrusher lazy initialization.""" + + def test_crusher_is_lazily_initialized(self): + """SmartCrusher is not created until first access.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + + # Directly check internal state - crusher should be None initially + assert hook._crusher is None + + # Access the crusher property + crusher = hook.crusher + + # Now it should be initialized + assert crusher is not None + assert hook._crusher is crusher + + def test_crusher_uses_configured_min_tokens(self): + """SmartCrusher uses min_tokens_to_compress from hook config.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider(min_tokens_to_compress=250) + + crusher = hook.crusher + + # The crusher config should have our min_tokens setting + assert crusher.config.min_tokens_to_crush == 250 + + +class TestTokenEstimation: + """Tests for _estimate_tokens helper method.""" + + def test_estimate_tokens_empty_string(self): + """Estimate returns 0 for empty string.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + assert hook._estimate_tokens("") == 0 + + def test_estimate_tokens_short_string(self): + """Estimate uses ~4 chars per token heuristic.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + + # 12 chars = 3 tokens (12 // 4) + assert hook._estimate_tokens("hello world!") == 3 + + # 20 chars = 5 tokens + assert hook._estimate_tokens("a" * 20) == 5 + + +class TestExtractTextContent: + """Tests for _extract_text_content helper method.""" + + def test_extract_from_text_content(self): + """Extract text from content with text field.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + result = {"content": [{"text": "Hello world"}]} + + extracted = hook._extract_text_content(result) + assert extracted == "Hello world" + + def test_extract_from_json_content(self): + """Extract and serialize JSON content.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + result = {"content": [{"json": {"key": "value"}}]} + + extracted = hook._extract_text_content(result) + assert extracted == '{"key": "value"}' + + def test_extract_empty_content(self): + """Return empty string for empty content.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + result = {"content": []} + + extracted = hook._extract_text_content(result) + assert extracted == "" + + def test_extract_missing_content(self): + """Return empty string for missing content key.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + result = {} + + extracted = hook._extract_text_content(result) + assert extracted == "" + + +class TestShouldSkipCompression: + """Tests for _should_skip_compression helper method.""" + + def test_skip_when_compression_disabled(self): + """Skip compression when compress_tool_outputs is False.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider(compress_tool_outputs=False) + result = {"content": [{"text": "data"}]} + + skip_reason = hook._should_skip_compression(result) + assert skip_reason == "compression_disabled" + + def test_skip_error_results_when_preserve_errors_true(self): + """Skip error results when preserve_errors is True.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider(preserve_errors=True) + result = {"status": "error", "content": [{"text": "Error message"}]} + + skip_reason = hook._should_skip_compression(result) + assert skip_reason == "error_result_preserved" + + def test_allow_error_results_when_preserve_errors_false(self): + """Allow error results when preserve_errors is False.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider(preserve_errors=False) + result = {"status": "error", "content": [{"text": "Error message"}]} + + skip_reason = hook._should_skip_compression(result) + assert skip_reason is None + + def test_skip_empty_content(self): + """Skip results with empty content.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + result = {"content": []} + + skip_reason = hook._should_skip_compression(result) + assert skip_reason == "empty_content" + + def test_allow_valid_content(self): + """Allow results with valid content.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + result = {"content": [{"text": "some data"}]} + + skip_reason = hook._should_skip_compression(result) + assert skip_reason is None + + +class TestCompressToolResult: + """Tests for _compress_tool_result hook handler.""" + + def test_compress_large_tool_output(self): + """Compresses large tool output and tracks metrics.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider( + compress_tool_outputs=True, + min_tokens_to_compress=10, # Low threshold for testing + ) + + # Create large JSON output (50 items) + large_data = [{"id": i, "value": f"item-{i}", "data": "x" * 50} for i in range(50)] + large_json = json.dumps(large_data) + + mock_event = MagicMock() + mock_event.tool_use = {"name": "get_items", "toolUseId": "tool-123"} + mock_event.result = {"content": [{"text": large_json}]} + + hook._compress_tool_result(mock_event) + + # Verify metrics were recorded + assert len(hook.metrics_history) == 1 + metrics = hook.metrics_history[0] + assert metrics.tool_name == "get_items" + assert metrics.tool_use_id == "tool-123" + assert metrics.tokens_before > 0 + + def test_skip_compression_below_threshold(self): + """Does not compress output below token threshold.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider( + compress_tool_outputs=True, + min_tokens_to_compress=10000, # High threshold + ) + + mock_event = MagicMock() + mock_event.tool_use = {"name": "small_tool", "toolUseId": "tool-456"} + mock_event.result = {"content": [{"text": '{"status": "ok"}'}]} + + hook._compress_tool_result(mock_event) + + # Metrics should show skipped compression + assert len(hook.metrics_history) == 1 + metrics = hook.metrics_history[0] + assert metrics.was_compressed is False + assert "below_threshold" in metrics.skip_reason + + def test_skip_compression_when_disabled(self): + """Does not compress when compression is disabled.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider(compress_tool_outputs=False) + + mock_event = MagicMock() + mock_event.tool_use = {"name": "test_tool", "toolUseId": "tool-789"} + mock_event.result = {"content": [{"text": '{"data": "value"}'}]} + + hook._compress_tool_result(mock_event) + + # Metrics should show compression disabled + assert len(hook.metrics_history) == 1 + metrics = hook.metrics_history[0] + assert metrics.was_compressed is False + assert metrics.skip_reason == "compression_disabled" + + +class TestMetricsTracking: + """Tests for metrics tracking and aggregation.""" + + def test_total_tokens_saved_accumulates(self): + """total_tokens_saved accumulates across compressions.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider( + compress_tool_outputs=True, + min_tokens_to_compress=10, + ) + + # Simulate two compressions with savings + for i in range(2): + large_data = [{"id": j, "data": "x" * 100} for j in range(50)] + mock_event = MagicMock() + mock_event.tool_use = {"name": f"tool_{i}", "toolUseId": f"id_{i}"} + mock_event.result = {"content": [{"text": json.dumps(large_data)}]} + + hook._compress_tool_result(mock_event) + + # Should have accumulated some savings + compressed_count = sum(1 for m in hook.metrics_history if m.was_compressed) + if compressed_count > 0: + assert hook.total_tokens_saved >= 0 + + def test_metrics_history_bounded_to_100(self): + """metrics_history keeps only last 100 entries.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider( + compress_tool_outputs=True, + min_tokens_to_compress=10, + ) + + # Directly add 150 metrics + for i in range(150): + hook._record_metrics( + request_id=f"req_{i}", + tool_name=f"tool_{i}", + tool_use_id=f"id_{i}", + tokens_before=100, + tokens_after=50, + was_compressed=True, + skip_reason=None, + ) + + # Should be bounded at 100 + assert len(hook.metrics_history) == 100 + + # Should contain the most recent entries + last_metric = hook.metrics_history[-1] + assert last_metric.request_id == "req_149" + + +class TestGetSavingsSummary: + """Tests for get_savings_summary method.""" + + def test_empty_summary(self): + """Returns zero values when no metrics recorded.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + summary = hook.get_savings_summary() + + assert summary["total_requests"] == 0 + assert summary["compressed_requests"] == 0 + assert summary["total_tokens_saved"] == 0 + assert summary["average_savings_percent"] == 0.0 + + def test_summary_with_compressions(self): + """Returns correct summary with recorded compressions.""" + from headroom.integrations.strands import HeadroomHookProvider + from headroom.integrations.strands.hooks import CompressionMetrics + + hook = HeadroomHookProvider() + + # Add metrics manually + hook._metrics_history = [ + CompressionMetrics( + request_id="1", + timestamp=datetime.now(timezone.utc), + tool_name="tool_a", + tool_use_id="id_1", + tokens_before=100, + tokens_after=60, + tokens_saved=40, + savings_percent=40.0, + was_compressed=True, + skip_reason=None, + ), + CompressionMetrics( + request_id="2", + timestamp=datetime.now(timezone.utc), + tool_name="tool_b", + tool_use_id="id_2", + tokens_before=200, + tokens_after=100, + tokens_saved=100, + savings_percent=50.0, + was_compressed=True, + skip_reason=None, + ), + CompressionMetrics( + request_id="3", + timestamp=datetime.now(timezone.utc), + tool_name="tool_c", + tool_use_id="id_3", + tokens_before=50, + tokens_after=50, + tokens_saved=0, + savings_percent=0.0, + was_compressed=False, + skip_reason="below_threshold", + ), + ] + hook._total_tokens_saved = 140 + + summary = hook.get_savings_summary() + + assert summary["total_requests"] == 3 + assert summary["compressed_requests"] == 2 + assert summary["total_tokens_saved"] == 140 + assert summary["average_savings_percent"] == 45.0 # (40 + 50) / 2 + assert summary["total_tokens_before"] == 350 + assert summary["total_tokens_after"] == 210 + + +class TestReset: + """Tests for reset method.""" + + def test_reset_clears_all_state(self): + """reset() clears all tracked state.""" + from headroom.integrations.strands import HeadroomHookProvider + from headroom.integrations.strands.hooks import CompressionMetrics + + hook = HeadroomHookProvider() + + # Add some state + hook._metrics_history = [ + CompressionMetrics( + request_id="1", + timestamp=datetime.now(timezone.utc), + tool_name="test", + tool_use_id="id_1", + tokens_before=100, + tokens_after=50, + tokens_saved=50, + savings_percent=50.0, + was_compressed=True, + ) + ] + hook._total_tokens_saved = 50 + + # Reset + hook.reset() + + # Verify all state cleared + assert hook._metrics_history == [] + assert hook._total_tokens_saved == 0 + assert hook.total_tokens_saved == 0 + assert len(hook.metrics_history) == 0 + + +class TestThreadSafety: + """Tests for thread-safety of metrics tracking.""" + + def test_concurrent_metric_recording(self): + """Metrics recording is thread-safe.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + + def record_metrics(thread_id): + for i in range(10): + hook._record_metrics( + request_id=f"thread_{thread_id}_req_{i}", + tool_name=f"tool_{thread_id}_{i}", + tool_use_id=f"id_{thread_id}_{i}", + tokens_before=100, + tokens_after=50, + was_compressed=True, + skip_reason=None, + ) + + threads = [] + for t_id in range(5): + t = threading.Thread(target=record_metrics, args=(t_id,)) + threads.append(t) + t.start() + + for t in threads: + t.join() + + # Should have recorded 50 metrics (5 threads * 10 each) + # But bounded to 100, so if we had more it would be truncated + assert len(hook.metrics_history) == 50 + assert hook.total_tokens_saved == 50 * 50 # 50 metrics * 50 tokens each + + +class TestUpdateResultContent: + """Tests for _update_result_content helper method.""" + + def test_update_preserves_json_structure(self): + """Updates preserve JSON structure when possible.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + result = {"content": [{"json": {"original": "data"}}]} + + compressed = '{"compressed": "data"}' + hook._update_result_content(result, compressed) + + # Should update with parsed JSON + assert result["content"] == [{"json": {"compressed": "data"}}] + + def test_update_uses_text_for_non_json(self): + """Updates use text format for non-JSON content.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + result = {"content": [{"text": "original text"}]} + + compressed = "compressed text" + hook._update_result_content(result, compressed) + + assert result["content"] == [{"text": "compressed text"}] + + def test_update_creates_content_if_empty(self): + """Creates content list if missing.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + result = {"content": []} + + hook._update_result_content(result, "new content") + + assert result["content"] == [{"text": "new content"}] diff --git a/tests/integrations/test_strands/test_model.py b/tests/integrations/test_strands/test_model.py new file mode 100644 index 000000000..0a8eae041 --- /dev/null +++ b/tests/integrations/test_strands/test_model.py @@ -0,0 +1,673 @@ +"""Real-world integration tests for Strands HeadroomStrandsModel. + +These tests use actual AWS Bedrock API calls with real credentials. +NO MOCKS - all tests hit the real Bedrock API. + +Skip in CI if AWS credentials are not available. +""" + +from __future__ import annotations + +import json +import os + +import pytest + +# Check for AWS credentials availability +SKIP_BEDROCK = not ( + os.environ.get("AWS_ACCESS_KEY_ID") + or os.environ.get("AWS_PROFILE") + or os.path.exists(os.path.expanduser("~/.aws/credentials")) +) + +# Check if strands-agents is installed +try: + from strands import Agent, tool + from strands.models import BedrockModel + + STRANDS_AVAILABLE = True +except ImportError: + STRANDS_AVAILABLE = False + + # Provide a no-op decorator when strands is not installed + def tool(fn): + return fn + + Agent = None # type: ignore + BedrockModel = None # type: ignore + +# Skip all tests if dependencies not available +pytestmark = [ + pytest.mark.skipif(SKIP_BEDROCK, reason="AWS credentials not available"), + pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed"), +] + + +# ============================================================================ +# Test Tools - Generate realistic data for optimization testing +# These are defined with @tool decorator for use when strands is installed. +# When strands is not installed, the no-op decorator ensures import succeeds. +# ============================================================================ + + +@tool +def get_database_records(table: str, limit: int = 50) -> str: + """Fetch records from a database table. Returns JSON array. + + Args: + table: Name of the database table + limit: Maximum records to return + + Returns: + JSON array of database records + """ + records = [ + { + "id": i, + "table": table, + "created_at": f"2024-01-{(i % 28) + 1:02d}T{10 + (i % 12):02d}:00:00Z", + "updated_at": f"2024-01-{(i % 28) + 1:02d}T{11 + (i % 12):02d}:00:00Z", + "status": ["active", "inactive", "pending", "archived"][i % 4], + "priority": ["low", "medium", "high", "critical"][i % 4], + "data": { + "field1": f"value_{i}_{table}", + "field2": i * 100, + "field3": i % 2 == 0, + "metadata": { + "source": "database", + "version": f"1.{i % 10}.0", + "tags": [f"tag_{j}" for j in range(i % 5 + 1)], + }, + }, + "metrics": { + "read_count": i * 10, + "write_count": i * 5, + "error_count": i % 3, + "latency_ms": 50 + (i * 7) % 200, + }, + } + for i in range(limit) + ] + return json.dumps(records, indent=2) + + +@tool +def get_large_logs(query: str, count: int = 200) -> str: + """Fetch verbose log data that should trigger compression. + + Args: + query: Search query for logs + count: Number of log entries to return + + Returns: + JSON array of detailed log entries + """ + logs = [ + { + "log_id": f"log_{i:08d}", + "timestamp": f"2024-01-{(i % 28) + 1:02d}T{10 + (i % 12):02d}:{i % 60:02d}:00Z", + "level": ["DEBUG", "INFO", "WARN", "ERROR"][i % 4], + "service": f"service_{i % 10}", + "message": f"Processing request for query '{query}' - step {i}", + "request_id": f"req_{i:012d}", + "trace_id": f"trace_{i:016x}", + "span_id": f"span_{i:08x}", + "user_id": f"user_{i % 100:04d}", + "session_id": f"sess_{i:010d}", + "metadata": { + "host": f"server-{i % 20:02d}.example.com", + "region": ["us-west-2", "us-east-1", "eu-west-1", "ap-southeast-1"][i % 4], + "instance_type": ["t3.micro", "t3.small", "t3.medium", "t3.large"][i % 4], + "container_id": f"container_{i:08x}", + "kubernetes_pod": f"pod-{i:06d}", + "kubernetes_namespace": "production", + }, + "metrics": { + "duration_ms": 50 + (i * 3) % 500, + "memory_mb": 128 + (i * 7) % 1024, + "cpu_percent": 5 + (i * 2) % 95, + "network_bytes_in": i * 1024, + "network_bytes_out": i * 512, + }, + "tags": ["env:prod", f"version:1.{i % 10}.0", "team:backend"], + } + for i in range(count) + ] + return json.dumps(logs, indent=2) + + +@tool +def analyze_metrics(metric_type: str) -> str: + """Analyze system metrics. Returns detailed metrics data. + + Args: + metric_type: Type of metrics to analyze (cpu, memory, network, disk) + + Returns: + JSON object with metric analysis + """ + data_points = [ + { + "timestamp": f"2024-01-15T{10 + (i % 12):02d}:{(i * 5) % 60:02d}:00Z", + "value": 20 + (i * 3) % 80, + "unit": {"cpu": "%", "memory": "MB", "network": "Mbps", "disk": "GB"}.get( + metric_type, "units" + ), + "host": f"server-{(i % 5) + 1:02d}", + "region": ["us-west-2", "us-east-1", "eu-west-1"][i % 3], + "metadata": { + "collection_interval": 60, + "aggregation": "avg", + "quality": "good" if i % 5 != 0 else "degraded", + }, + } + for i in range(100) + ] + + return json.dumps( + { + "metric_type": metric_type, + "time_range": {"start": "2024-01-15T10:00:00Z", "end": "2024-01-15T22:00:00Z"}, + "data_points": data_points, + "summary": { + "min": 20, + "max": 99, + "avg": 55.5, + "p50": 52, + "p95": 90, + "p99": 97, + }, + }, + indent=2, + ) + + +@tool +def quick_lookup(key: str) -> str: + """Quick key-value lookup. Returns small response. + + Args: + key: The key to look up + + Returns: + Small JSON with the value + """ + return json.dumps({"key": key, "value": f"result_for_{key}", "found": True}) + + +@tool +def math_operation(x: float, y: float, op: str) -> str: + """Perform a math operation. + + Args: + x: First operand + y: Second operand + op: Operation (add, sub, mul, div) + + Returns: + Result of the operation + """ + operations = { + "add": x + y, + "sub": x - y, + "mul": x * y, + "div": x / y if y != 0 else None, + } + result = operations.get(op, None) + return json.dumps({"x": x, "y": y, "operation": op, "result": result}) + + +# ============================================================================ +# Test Class for HeadroomStrandsModel +# ============================================================================ + + +@pytest.mark.skipif(SKIP_BEDROCK, reason="AWS credentials not available") +@pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed") +class TestHeadroomStrandsModelReal: + """Real-world integration tests for HeadroomStrandsModel with Bedrock.""" + + @pytest.fixture + def base_bedrock_model(self): + """Create a base BedrockModel instance using Claude 3 Haiku (fast and cheap).""" + return BedrockModel( + model_id="anthropic.claude-3-haiku-20240307-v1:0", + region_name="us-west-2", + temperature=0.1, + ) + + @pytest.fixture + def wrapped_model(self, base_bedrock_model): + """Create a HeadroomStrandsModel wrapping the Bedrock model.""" + from headroom.integrations.strands import HeadroomStrandsModel + + return HeadroomStrandsModel( + wrapped_model=base_bedrock_model, + auto_detect_provider=True, + ) + + def test_stream_returns_proper_events(self, wrapped_model): + """Test that stream() works and returns proper StreamEvents. + + The Strands Agent uses the model's stream() method internally. + This test verifies that the wrapped model properly streams responses. + """ + wrapped_model.reset() + + agent = Agent(model=wrapped_model) + + # Make a request - the agent internally calls stream() on the model + result = agent("Count from 1 to 5, one number per line.") + + # Verify we got a response (proves streaming worked) + assert result is not None + response_text = str(result) + assert len(response_text) > 0 + + # The response should contain numbers 1-5 + for num in ["1", "2", "3", "4", "5"]: + assert num in response_text, f"Expected {num} in response" + + # Metrics should be tracked (proves stream() was intercepted properly) + metrics = wrapped_model.get_savings_summary() + assert metrics["total_requests"] >= 1, "stream() should track requests" + + def test_messages_optimized_large_conversations(self, wrapped_model): + """Test that messages are actually optimized (tokens_before > tokens_after for large conversations). + + This test builds up a large conversation context through tool calls + with verbose JSON responses, then verifies that optimization occurs. + """ + wrapped_model.reset() + + agent = Agent(model=wrapped_model, tools=[get_large_logs, get_database_records]) + + # First request - get large logs (200 entries with verbose data) + agent( + "Search for logs containing 'error' and get 200 entries using get_large_logs. " + "Tell me how many ERROR level logs there are." + ) + + # Second request - more tool output, context grows + agent( + "Now get 100 records from the 'events' table using get_database_records. " + "How many records have 'active' status?" + ) + + # Third request - even more context + agent( + "Based on all the data you've seen, give me a one-sentence summary " + "of the system health." + ) + + # Check optimization metrics + metrics = wrapped_model.get_savings_summary() + + # Should have processed multiple requests + assert metrics["total_requests"] >= 1, "Should have processed requests" + + # With large tool outputs, tokens_before should be significant + assert metrics["total_tokens_before"] > 0, "Should have counted input tokens" + + # The key assertion: optimization should reduce tokens + # (tokens_before >= tokens_after, with strict > when there's compressible content) + assert metrics["total_tokens_before"] >= metrics["total_tokens_after"], ( + f"Optimization should not increase tokens: " + f"before={metrics['total_tokens_before']}, after={metrics['total_tokens_after']}" + ) + + # Check history shows optimization was tracked + history = wrapped_model.metrics_history + assert len(history) >= 1, "Should have metrics history" + + # Verify individual requests track before/after properly + for m in history: + assert m.tokens_before >= m.tokens_after, ( + f"Each request should have tokens_before >= tokens_after: " + f"request_id={m.request_id}, before={m.tokens_before}, after={m.tokens_after}" + ) + + def test_get_savings_summary_returns_correct_metrics(self, wrapped_model): + """Test that get_savings_summary() returns correct metrics. + + Verifies the structure and accuracy of the savings summary. + """ + wrapped_model.reset() + + agent = Agent(model=wrapped_model, tools=[get_database_records]) + + # Make a few requests + agent("Get 30 records from 'users' table.") + agent("Get 30 records from 'orders' table.") + + # Get the summary + summary = wrapped_model.get_savings_summary() + + # Verify required keys exist + required_keys = [ + "total_requests", + "total_tokens_saved", + "average_savings_percent", + "total_tokens_before", + "total_tokens_after", + ] + for key in required_keys: + assert key in summary, f"Summary missing required key: {key}" + + # Verify values are sensible + assert summary["total_requests"] >= 1, "Should have at least one request" + assert summary["total_tokens_before"] >= 0, "tokens_before should be non-negative" + assert summary["total_tokens_after"] >= 0, "tokens_after should be non-negative" + assert summary["total_tokens_saved"] >= 0, "tokens_saved should be non-negative" + assert 0 <= summary["average_savings_percent"] <= 100, ( + "average_savings_percent should be between 0 and 100" + ) + + # Verify mathematical consistency + expected_saved = summary["total_tokens_before"] - summary["total_tokens_after"] + assert summary["total_tokens_saved"] == expected_saved, ( + f"tokens_saved should equal tokens_before - tokens_after: " + f"saved={summary['total_tokens_saved']}, expected={expected_saved}" + ) + + def test_reset_clears_all_metrics(self, wrapped_model): + """Test that reset() clears all accumulated metrics. + + Verifies that reset() properly clears: + - total_tokens_saved + - metrics_history + - The summary returned by get_savings_summary() + """ + # Make some requests to accumulate metrics + agent = Agent(model=wrapped_model) + agent("Say 'hello world'") + agent("Say 'goodbye world'") + + # Verify we have metrics before reset + assert wrapped_model.total_tokens_saved >= 0 + pre_reset_requests = wrapped_model.get_savings_summary()["total_requests"] + assert pre_reset_requests >= 1, "Should have requests before reset" + + # Call reset + wrapped_model.reset() + + # Verify all metrics are cleared + assert wrapped_model.total_tokens_saved == 0, "total_tokens_saved should be 0 after reset" + assert len(wrapped_model.metrics_history) == 0, ( + "metrics_history should be empty after reset" + ) + + # Verify get_savings_summary reflects the reset + summary = wrapped_model.get_savings_summary() + assert summary["total_requests"] == 0, "total_requests should be 0 after reset" + assert summary["total_tokens_saved"] == 0, "total_tokens_saved should be 0 after reset" + assert summary["total_tokens_before"] == 0, "total_tokens_before should be 0 after reset" + assert summary["total_tokens_after"] == 0, "total_tokens_after should be 0 after reset" + + # Verify we can still make requests after reset + agent = Agent(model=wrapped_model) + agent("Say 'post-reset test'") + + post_reset_summary = wrapped_model.get_savings_summary() + assert post_reset_summary["total_requests"] >= 1, "Should track requests after reset" + + def test_model_wrapper_basic_response(self, wrapped_model): + """Test that wrapped model produces valid responses.""" + agent = Agent(model=wrapped_model) + + result = agent("Say 'Hello, Headroom!' and nothing else.") + + assert result is not None + content = str(result) + assert len(content) > 0 + + def test_model_wrapper_with_tools(self, wrapped_model): + """Test that wrapped model works correctly with tools.""" + wrapped_model.reset() + + agent = Agent(model=wrapped_model, tools=[quick_lookup, math_operation, analyze_metrics]) + + result = agent( + "Please do these tasks: " + "1. Look up the key 'config_setting' using quick_lookup. " + "2. Calculate 15.5 multiplied by 4 using math_operation. " + "3. Tell me the results." + ) + + assert result is not None + + metrics = wrapped_model.get_savings_summary() + assert metrics["total_requests"] >= 1 + + def test_model_wrapper_metrics_tracking(self, wrapped_model): + """Test that metrics are accurately tracked across requests.""" + wrapped_model.reset() + + agent = Agent(model=wrapped_model, tools=[get_database_records]) + + # Make several requests + agent("Get 20 records from 'products' table.") + agent("Get 20 records from 'customers' table.") + agent("Summarize both sets of records.") + + metrics = wrapped_model.get_savings_summary() + + assert metrics["total_requests"] >= 1 + assert metrics["total_tokens_before"] >= metrics["total_tokens_after"] + + if metrics["total_tokens_saved"] > 0: + assert metrics["average_savings_percent"] >= 0 + assert metrics["average_savings_percent"] <= 100 + + # History should be bounded + assert len(wrapped_model.metrics_history) <= 100 + + def test_model_wrapper_attribute_forwarding(self, base_bedrock_model): + """Test that attributes are forwarded to wrapped model.""" + from headroom.integrations.strands import HeadroomStrandsModel + + wrapped = HeadroomStrandsModel( + wrapped_model=base_bedrock_model, + auto_detect_provider=True, + ) + + # The wrapper should forward config to the wrapped model (Strands stores model_id in config) + assert hasattr(wrapped, "config") + config = wrapped.config + assert isinstance(config, dict) + assert "model_id" in config + + # Access wrapped model directly + assert wrapped.wrapped_model is base_bedrock_model + + def test_model_wrapper_custom_config(self, base_bedrock_model): + """Test that custom HeadroomConfig is applied.""" + from headroom import HeadroomConfig + from headroom.integrations.strands import HeadroomStrandsModel + + custom_config = HeadroomConfig() + custom_config.smart_crusher.min_tokens_to_crush = 50 + custom_config.smart_crusher.max_items_after_crush = 10 + + wrapped = HeadroomStrandsModel( + wrapped_model=base_bedrock_model, + config=custom_config, + auto_detect_provider=True, + ) + + assert wrapped.headroom_config is custom_config + assert wrapped.headroom_config.smart_crusher.min_tokens_to_crush == 50 + + # The model should still work + agent = Agent(model=wrapped) + result = agent("Say 'test'") + assert result is not None + + def test_model_wrapper_provider_detection(self, base_bedrock_model): + """Test that provider is auto-detected correctly for Bedrock Claude.""" + from headroom.integrations.strands import HeadroomStrandsModel + from headroom.providers import AnthropicProvider + + wrapped = HeadroomStrandsModel( + wrapped_model=base_bedrock_model, + auto_detect_provider=True, + ) + + # Access pipeline to trigger lazy initialization + _ = wrapped.pipeline + + # For Bedrock Claude models, should detect Anthropic provider + assert wrapped._headroom_provider is not None + assert isinstance(wrapped._headroom_provider, AnthropicProvider) + + def test_model_wrapper_handles_large_context(self, wrapped_model): + """Test that wrapper handles large context appropriately.""" + wrapped_model.reset() + + agent = Agent(model=wrapped_model, tools=[analyze_metrics, get_database_records]) + + # Build up context with large tool outputs + agent("Analyze CPU metrics using analyze_metrics.") + agent("Get 50 records from 'logs' table using get_database_records.") + agent("Based on everything, what patterns do you see?") + + metrics = wrapped_model.get_savings_summary() + assert metrics["total_requests"] >= 1 + assert metrics["total_tokens_before"] > 0 + + def test_model_wrapper_empty_messages(self, base_bedrock_model): + """Test that wrapper handles edge cases gracefully.""" + from headroom.integrations.strands import HeadroomStrandsModel + + wrapped = HeadroomStrandsModel( + wrapped_model=base_bedrock_model, + auto_detect_provider=True, + ) + + # Test with minimal input + agent = Agent(model=wrapped) + result = agent("Hi") + + assert result is not None + + def test_model_wrapper_thread_safety(self, base_bedrock_model): + """Test that wrapper is thread-safe for metrics tracking.""" + import threading + import time + + from headroom.integrations.strands import HeadroomStrandsModel + + wrapped = HeadroomStrandsModel( + wrapped_model=base_bedrock_model, + auto_detect_provider=True, + ) + + agent = Agent(model=wrapped) + + results = [] + errors = [] + + def make_request(msg: str): + try: + result = agent(msg) + results.append(result) + except Exception as e: + errors.append(e) + + threads = [] + messages = ["Say 'one'", "Say 'two'", "Say 'three'"] + + for msg in messages: + t = threading.Thread(target=make_request, args=(msg,)) + threads.append(t) + t.start() + time.sleep(0.5) # Small delay to avoid rate limiting + + for t in threads: + t.join(timeout=60) + + # Should have some results (may have errors due to rate limiting) + assert len(results) > 0 or len(errors) > 0 + + # Metrics should be consistent + metrics = wrapped.get_savings_summary() + assert metrics["total_tokens_before"] >= metrics["total_tokens_after"] + + +# ============================================================================ +# Test Class for optimize_messages standalone function +# ============================================================================ + + +@pytest.mark.skipif(SKIP_BEDROCK, reason="AWS credentials not available") +@pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed") +class TestOptimizeMessagesFunction: + """Tests for the standalone optimize_messages function.""" + + def test_optimize_messages_basic(self): + """Test basic message optimization.""" + from headroom.integrations.strands import optimize_messages + + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello!"}, + {"role": "assistant", "content": "Hi there! How can I help you today?"}, + ] + + optimized, metrics = optimize_messages(messages) + + assert len(optimized) > 0 + + assert "tokens_before" in metrics + assert "tokens_after" in metrics + assert "tokens_saved" in metrics + assert metrics["tokens_before"] >= 0 + assert metrics["tokens_after"] >= 0 + + def test_optimize_messages_with_tool_content(self): + """Test optimization of messages containing tool responses.""" + from headroom.integrations.strands import optimize_messages + + # Create messages with large tool output + large_data = json.dumps([{"id": i, "data": f"value_{i}" * 10} for i in range(100)]) + + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Get the data"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": {"name": "get_data", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "content": large_data, "tool_call_id": "call_123"}, + {"role": "assistant", "content": "Here is the data summary..."}, + ] + + optimized, metrics = optimize_messages(messages) + + assert len(optimized) > 0 + assert metrics["tokens_before"] >= 0 + + def test_optimize_messages_custom_config(self): + """Test optimization with custom config.""" + from headroom import HeadroomConfig + from headroom.integrations.strands import optimize_messages + + config = HeadroomConfig() + config.smart_crusher.enabled = True + config.smart_crusher.min_tokens_to_crush = 10 + + messages = [ + {"role": "user", "content": "Hello!"}, + ] + + optimized, metrics = optimize_messages(messages, config=config) + + assert len(optimized) > 0 + assert "tokens_before" in metrics diff --git a/tests/integrations/test_strands/test_model_unit.py b/tests/integrations/test_strands/test_model_unit.py new file mode 100644 index 000000000..9848530bc --- /dev/null +++ b/tests/integrations/test_strands/test_model_unit.py @@ -0,0 +1,645 @@ +"""Unit tests for Strands HeadroomStrandsModel. + +These tests use mocks and do NOT require AWS credentials or strands-agents. +They test the internal logic of HeadroomStrandsModel in isolation. + +For real integration tests, see test_model.py. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch + +import pytest + +# Check if strands-agents is installed for proper skip handling +try: + import strands # noqa: F401 + + STRANDS_AVAILABLE = True +except ImportError: + STRANDS_AVAILABLE = False + + +# Skip all tests if Strands not installed +pytestmark = pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed") + + +# ============================================================================ +# Fixtures +# ============================================================================ + + +@pytest.fixture +def mock_strands_model(): + """Create a mock Strands model.""" + mock = MagicMock() + mock.config = {"model_id": "anthropic.claude-3-haiku-20240307-v1:0"} + mock.get_config.return_value = mock.config + + # Mock the stream method as an async generator + async def mock_stream(*args, **kwargs): + yield {"type": "content", "data": "Hello"} + yield {"type": "content", "data": " world"} + yield {"type": "stop"} + + mock.stream = mock_stream + return mock + + +@pytest.fixture +def sample_messages(): + """Sample messages in Strands/OpenAI format.""" + return [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is the capital of France?"}, + ] + + +@pytest.fixture +def large_conversation(): + """Large conversation with many turns for compression testing.""" + messages = [{"role": "system", "content": "You are a helpful assistant."}] + for i in range(50): + messages.append({"role": "user", "content": f"Question {i}: What is {i} + {i}?"}) + messages.append({"role": "assistant", "content": f"The answer is {i + i}."}) + return messages + + +# ============================================================================ +# Test Classes +# ============================================================================ + + +class TestHeadroomStrandsModelInit: + """Tests for HeadroomStrandsModel initialization.""" + + def test_init_with_defaults(self, mock_strands_model): + """Initialize with default settings.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + assert model.wrapped_model is mock_strands_model + assert model.total_tokens_saved == 0 + assert model.metrics_history == [] + assert model.auto_detect_provider is True + + def test_init_with_custom_config(self, mock_strands_model): + """Initialize with custom HeadroomConfig.""" + from headroom import HeadroomConfig + from headroom.integrations.strands import HeadroomStrandsModel + + config = HeadroomConfig() + config.smart_crusher.min_tokens_to_crush = 100 + + model = HeadroomStrandsModel( + wrapped_model=mock_strands_model, + config=config, + auto_detect_provider=False, + ) + + assert model.headroom_config is config + assert model.auto_detect_provider is False + + def test_init_requires_wrapped_model(self): + """Raises ValueError if wrapped_model is None.""" + from headroom.integrations.strands import HeadroomStrandsModel + + with pytest.raises(ValueError, match="wrapped_model cannot be None"): + HeadroomStrandsModel(wrapped_model=None) + + +class TestAttributeForwarding: + """Tests for attribute forwarding to wrapped model.""" + + def test_forwards_unknown_attributes(self, mock_strands_model): + """Forwards unknown attributes to wrapped model.""" + from headroom.integrations.strands import HeadroomStrandsModel + + mock_strands_model.custom_attr = "custom_value" + mock_strands_model.another_attr = 42 + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + assert model.custom_attr == "custom_value" + assert model.another_attr == 42 + + def test_forwards_config_property(self, mock_strands_model): + """Forwards config property to wrapped model.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + config = model.config + assert config is mock_strands_model.config + + def test_does_not_forward_internal_attrs(self, mock_strands_model): + """Does not forward internal wrapper attributes.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + # These should be wrapper's own attributes + assert model.wrapped_model is mock_strands_model + assert model.total_tokens_saved == 0 + assert model.metrics_history == [] + + def test_get_config_delegates(self, mock_strands_model): + """get_config() delegates to wrapped model.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + config = model.get_config() + assert config == mock_strands_model.get_config() + + def test_update_config_delegates(self, mock_strands_model): + """update_config() delegates to wrapped model.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + model.update_config(temperature=0.5) + mock_strands_model.update_config.assert_called_once_with(temperature=0.5) + + +class TestMessageConversion: + """Tests for message format conversion.""" + + def test_convert_dict_messages(self, mock_strands_model, sample_messages): + """Converts dict messages to OpenAI format.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + converted = model._convert_messages_to_openai(sample_messages) + + assert len(converted) == 2 + assert converted[0]["role"] == "system" + assert converted[0]["content"] == "You are a helpful assistant." + assert converted[1]["role"] == "user" + assert converted[1]["content"] == "What is the capital of France?" + + def test_convert_messages_with_tool_calls(self, mock_strands_model): + """Converts messages with tool calls.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + messages = [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_123", "type": "function", "function": {"name": "search"}} + ], + }, + { + "role": "tool", + "content": '{"results": []}', + "tool_call_id": "call_123", + "name": "search", + }, + ] + + converted = model._convert_messages_to_openai(messages) + + assert len(converted) == 2 + assert "tool_calls" in converted[0] + assert converted[1]["tool_call_id"] == "call_123" + assert converted[1]["name"] == "search" + + def test_convert_message_objects(self, mock_strands_model): + """Converts message objects with role/content attributes.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + # Create mock message objects + msg1 = MagicMock() + msg1.role = "user" + msg1.content = "Hello" + msg1.tool_calls = None + msg1.tool_call_id = None + msg1.name = None + + msg2 = MagicMock() + msg2.role = "assistant" + msg2.content = "Hi there!" + msg2.tool_calls = None + msg2.tool_call_id = None + msg2.name = None + + converted = model._convert_messages_to_openai([msg1, msg2]) + + assert len(converted) == 2 + assert converted[0]["role"] == "user" + assert converted[0]["content"] == "Hello" + assert converted[1]["role"] == "assistant" + assert converted[1]["content"] == "Hi there!" + + def test_convert_handles_content_list(self, mock_strands_model): + """Converts messages with content as list (content blocks).""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Look at this:"}, + {"type": "image", "source": {"data": "base64..."}}, + ], + } + ] + + converted = model._convert_messages_to_openai(messages) + + assert len(converted) == 1 + assert isinstance(converted[0]["content"], list) + assert len(converted[0]["content"]) == 2 + + +class TestOptimizeMessages: + """Tests for _optimize_messages method.""" + + def test_optimize_returns_metrics(self, mock_strands_model, sample_messages): + """_optimize_messages returns messages and metrics.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + # Mock the pipeline by setting _pipeline directly and mocking _headroom_provider + mock_pipeline = MagicMock() + mock_result = MagicMock() + mock_result.messages = sample_messages + mock_result.tokens_before = 50 + mock_result.tokens_after = 40 + mock_result.transforms_applied = ["cache_aligner"] + mock_pipeline.apply.return_value = mock_result + + model._pipeline = mock_pipeline + model._headroom_provider = MagicMock() + model._headroom_provider.get_context_limit.return_value = 128000 + + optimized, metrics = model._optimize_messages(sample_messages) + + assert len(optimized) == 2 + assert metrics.tokens_before == 50 + assert metrics.tokens_after == 40 + assert metrics.tokens_saved == 10 + assert "cache_aligner" in metrics.transforms_applied + + def test_optimize_handles_empty_messages(self, mock_strands_model): + """_optimize_messages handles empty message list.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + optimized, metrics = model._optimize_messages([]) + + assert optimized == [] + assert metrics.tokens_before == 0 + assert metrics.tokens_after == 0 + assert metrics.tokens_saved == 0 + + def test_optimize_tracks_metrics(self, mock_strands_model, sample_messages): + """_optimize_messages tracks metrics in history.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + # Mock the pipeline by setting _pipeline directly + mock_pipeline = MagicMock() + mock_result = MagicMock() + mock_result.messages = sample_messages + mock_result.tokens_before = 100 + mock_result.tokens_after = 80 + mock_result.transforms_applied = [] + mock_pipeline.apply.return_value = mock_result + + model._pipeline = mock_pipeline + model._headroom_provider = MagicMock() + model._headroom_provider.get_context_limit.return_value = 128000 + + model._optimize_messages(sample_messages) + + assert len(model.metrics_history) == 1 + assert model.metrics_history[0].tokens_saved == 20 + assert model.total_tokens_saved == 20 + + def test_optimize_handles_pipeline_errors(self, mock_strands_model, sample_messages): + """_optimize_messages falls back on pipeline errors.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + # Mock the pipeline to raise an error + mock_pipeline = MagicMock() + mock_pipeline.apply.side_effect = ValueError("Pipeline error") + + model._pipeline = mock_pipeline + model._headroom_provider = MagicMock() + model._headroom_provider.get_context_limit.return_value = 128000 + + # Should not raise, should fall back + optimized, metrics = model._optimize_messages(sample_messages) + + assert len(optimized) == len(sample_messages) + assert "fallback:error" in metrics.transforms_applied + + +class TestPipelineLazyInit: + """Tests for TransformPipeline lazy initialization.""" + + def test_pipeline_is_lazily_initialized(self, mock_strands_model): + """Pipeline is not created until first access.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + # Should be None initially + assert model._pipeline is None + + # Access pipeline property + with patch("headroom.integrations.strands.model.TransformPipeline"): + _ = model.pipeline + + # Now should be initialized + assert model._pipeline is not None + + +class TestGetSavingsSummary: + """Tests for get_savings_summary method.""" + + def test_empty_summary(self, mock_strands_model): + """Returns zero values when no metrics recorded.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + summary = model.get_savings_summary() + + assert summary["total_requests"] == 0 + assert summary["total_tokens_saved"] == 0 + assert summary["average_savings_percent"] == 0 + + def test_summary_with_metrics(self, mock_strands_model): + """Returns correct summary with recorded metrics.""" + from headroom.integrations.strands import HeadroomStrandsModel + from headroom.integrations.strands.model import OptimizationMetrics + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + # Add metrics manually + model._metrics_history = [ + OptimizationMetrics( + request_id="1", + timestamp=datetime.now(timezone.utc), + tokens_before=100, + tokens_after=80, + tokens_saved=20, + savings_percent=20.0, + transforms_applied=[], + model="test-model", + ), + OptimizationMetrics( + request_id="2", + timestamp=datetime.now(timezone.utc), + tokens_before=200, + tokens_after=120, + tokens_saved=80, + savings_percent=40.0, + transforms_applied=[], + model="test-model", + ), + ] + model._total_tokens_saved = 100 + + summary = model.get_savings_summary() + + assert summary["total_requests"] == 2 + assert summary["total_tokens_saved"] == 100 + assert summary["average_savings_percent"] == 30.0 # (20 + 40) / 2 + assert summary["total_tokens_before"] == 300 + assert summary["total_tokens_after"] == 200 + + +class TestReset: + """Tests for reset method.""" + + def test_reset_clears_all_state(self, mock_strands_model): + """reset() clears all tracked state.""" + from headroom.integrations.strands import HeadroomStrandsModel + from headroom.integrations.strands.model import OptimizationMetrics + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + # Add some state + model._metrics_history = [ + OptimizationMetrics( + request_id="1", + timestamp=datetime.now(timezone.utc), + tokens_before=100, + tokens_after=50, + tokens_saved=50, + savings_percent=50.0, + transforms_applied=[], + model="test", + ) + ] + model._total_tokens_saved = 50 + + # Reset + model.reset() + + # Verify all state cleared + assert model._metrics_history == [] + assert model._total_tokens_saved == 0 + assert model.total_tokens_saved == 0 + assert len(model.metrics_history) == 0 + + # Summary should reflect reset + summary = model.get_savings_summary() + assert summary["total_requests"] == 0 + + +class TestMetricsHistoryBound: + """Tests for metrics history bounding.""" + + def test_metrics_bounded_to_100(self, mock_strands_model): + """Metrics history is bounded to 100 entries.""" + from headroom.integrations.strands import HeadroomStrandsModel + from headroom.integrations.strands.model import OptimizationMetrics + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + # Add 150 metrics + for i in range(150): + model._metrics_history.append( + OptimizationMetrics( + request_id=f"req_{i}", + timestamp=datetime.now(timezone.utc), + tokens_before=100, + tokens_after=80, + tokens_saved=20, + savings_percent=20.0, + transforms_applied=[], + model="test", + ) + ) + # Simulate what _optimize_messages does + if len(model._metrics_history) > 100: + model._metrics_history = model._metrics_history[-100:] + + # Should be bounded at 100 + assert len(model.metrics_history) == 100 + + # Should contain the most recent entries + assert model.metrics_history[-1].request_id == "req_149" + + +class TestOptimizeMessagesFunction: + """Tests for standalone optimize_messages function.""" + + def test_optimize_messages_basic(self): + """optimize_messages processes messages and returns metrics.""" + from headroom.integrations.strands import optimize_messages + + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there!"}, + ] + + with patch("headroom.integrations.strands.model.TransformPipeline") as MockPipeline: + mock_instance = MagicMock() + mock_result = MagicMock() + mock_result.messages = messages + mock_result.tokens_before = 20 + mock_result.tokens_after = 15 + mock_result.transforms_applied = ["cache_aligner"] + mock_instance.apply.return_value = mock_result + MockPipeline.return_value = mock_instance + + optimized, metrics = optimize_messages(messages) + + assert len(optimized) == 2 + assert metrics["tokens_saved"] == 5 + assert metrics["savings_percent"] == 25.0 + + def test_optimize_messages_with_custom_config(self): + """optimize_messages uses custom config.""" + from headroom import HeadroomConfig + from headroom.integrations.strands import optimize_messages + + config = HeadroomConfig() + messages = [{"role": "user", "content": "Test"}] + + with patch("headroom.integrations.strands.model.TransformPipeline") as MockPipeline: + mock_instance = MagicMock() + mock_result = MagicMock() + mock_result.messages = messages + mock_result.tokens_before = 10 + mock_result.tokens_after = 10 + mock_result.transforms_applied = [] + mock_instance.apply.return_value = mock_result + MockPipeline.return_value = mock_instance + + optimized, metrics = optimize_messages(messages, config=config) + + # Verify config was passed to pipeline + MockPipeline.assert_called_once() + call_kwargs = MockPipeline.call_args[1] + assert call_kwargs["config"] is config + + +class TestStreamMethod: + """Tests for stream method.""" + + @pytest.mark.asyncio + async def test_stream_optimizes_messages(self, mock_strands_model, sample_messages): + """stream() applies optimization before calling wrapped model.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel(wrapped_model=mock_strands_model) + + # Mock the optimization + with patch.object(model, "_optimize_messages") as mock_optimize: + mock_optimize.return_value = ( + sample_messages, + MagicMock( + tokens_before=50, + tokens_after=40, + savings_percent=20.0, + ), + ) + + # Consume the stream + events = [] + async for event in model.stream(sample_messages): + events.append(event) + + # Should have called optimization + mock_optimize.assert_called_once() + + # Should have yielded events from wrapped model + assert len(events) > 0 + + +class TestStrandsAvailableFunction: + """Tests for strands_available function.""" + + def test_strands_available_returns_bool(self): + """strands_available() returns boolean.""" + from headroom.integrations.strands import strands_available + + result = strands_available() + + # Since we're in a test where strands is available (skipif passed) + assert isinstance(result, bool) + assert result is True + + +class TestRealHeadroomIntegration: + """Integration tests with real Headroom (no mocking).""" + + def test_real_optimization_with_mock_model(self, mock_strands_model, sample_messages): + """Test with real Headroom transforms (no API calls).""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel( + wrapped_model=mock_strands_model, + auto_detect_provider=False, # Use default OpenAI provider + ) + + # This calls real Headroom optimization + optimized, metrics = model._optimize_messages(sample_messages) + + # Should return valid messages + assert len(optimized) >= 1 + assert all("role" in m and "content" in m for m in optimized) + + # Metrics should be tracked + assert len(model.metrics_history) == 1 + assert metrics.tokens_before >= 0 + assert metrics.tokens_after >= 0 + + def test_large_conversation_handling(self, mock_strands_model, large_conversation): + """Large conversations are processed without errors.""" + from headroom.integrations.strands import HeadroomStrandsModel + + model = HeadroomStrandsModel( + wrapped_model=mock_strands_model, + auto_detect_provider=False, + ) + + # Should handle large conversation without errors + optimized, metrics = model._optimize_messages(large_conversation) + + # Should return messages + assert len(optimized) >= 1 + + # Metrics should show processing occurred + assert metrics.tokens_before > 0