akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #16 from chopratejas/feature/strands-integration

Browse source 
Feature/strands integration
This commit is contained in:
Tejas Chopra 2026-01-31 00:33:01 -08:00 • committed by GitHub
commit 986adae8ea
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 5433 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

13
README.md
View file

@ -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

56
examples/README.md
View file

@ -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/<example_name>.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.

1001
examples/strands_bedrock_demo.py Normal file
View file

File diff suppressed because it is too large Load diff

12
headroom/dashboard/__init__.py Normal file
View file

@ -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()

377
headroom/dashboard/templates/dashboard.html Normal file
View file

@ -0,0 +1,377 @@
<!DOCTYPE html>
<html lang="en" class="dark">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Headroom Dashboard</title>
<script src="https://cdn.tailwindcss.com"></script>
<script src="https://unpkg.com/htmx.org@1.9.10"></script>
<script src="https://unpkg.com/alpinejs@3.13.3/dist/cdn.min.js" defer></script>
<script>
tailwind.config = {
darkMode: 'class',
theme: {
extend: {
colors: {
surface: '#1a1a1a',
border: '#2a2a2a',
accent: '#22d3ee',
}
}
}
}
</script>
<style>
body { background: #0f0f0f; }
.sparkline { stroke: #22d3ee; stroke-width: 1.5; fill: none; }
.sparkline-area { fill: url(#sparkline-gradient); }
@keyframes pulse-subtle { 0%, 100% { opacity: 1; } 50% { opacity: 0.7; } }
.pulse-live { animation: pulse-subtle 2s ease-in-out infinite; }
</style>
</head>
<body class="text-gray-200 min-h-screen" x-data="dashboard()" x-init="init()">
<!-- Header -->
<header class="border-b border-border px-6 py-4 flex items-center justify-between">
<div class="flex items-center gap-4">
<h1 class="text-xl font-semibold tracking-tight">HEADROOM</h1>
<span class="text-xs text-gray-500 font-mono" x-text="'v' + version"></span>
</div>
<div class="flex items-center gap-6">
<div class="flex items-center gap-2">
<span class="text-xs text-gray-500">Status</span>
<span class="flex items-center gap-1.5">
<span class="w-2 h-2 rounded-full pulse-live"
:class="healthy ? 'bg-emerald-400' : 'bg-red-400'"></span>
<span class="text-sm" x-text="healthy ? 'Healthy' : 'Error'"></span>
</span>
</div>
<div class="text-xs text-gray-500">
Updated <span x-text="lastUpdate"></span>
</div>
</div>
</header>
<main class="p-6 max-w-7xl mx-auto">
<!-- Hero Metrics -->
<div class="grid grid-cols-2 md:grid-cols-4 gap-4 mb-6">
<!-- Requests -->
<div class="bg-surface rounded-lg p-4 border border-border">
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">Requests</div>
<div class="text-3xl font-light tabular-nums" x-text="formatNumber(stats.requests?.total || 0)"></div>
<div class="mt-2 h-8">
<svg class="w-full h-full" viewBox="0 0 100 32" preserveAspectRatio="none">
<defs>
<linearGradient id="sparkline-gradient" x1="0%" y1="0%" x2="0%" y2="100%">
<stop offset="0%" style="stop-color:#22d3ee;stop-opacity:0.3"/>
<stop offset="100%" style="stop-color:#22d3ee;stop-opacity:0"/>
</linearGradient>
</defs>
<path class="sparkline-area" :d="getSparklineArea(requestHistory)"></path>
<path class="sparkline" :d="getSparkline(requestHistory)"></path>
</svg>
</div>
</div>
<!-- Tokens Saved -->
<div class="bg-surface rounded-lg p-4 border border-border">
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">Tokens Saved</div>
<div class="flex items-baseline gap-2">
<span class="text-3xl font-light tabular-nums text-accent" x-text="formatNumber(stats.tokens?.saved || 0)"></span>
<span class="text-sm text-accent" x-text="(stats.tokens?.savings_percent || 0).toFixed(1) + '%'"></span>
</div>
<div class="mt-2 h-8">
<svg class="w-full h-full" viewBox="0 0 100 32" preserveAspectRatio="none">
<path class="sparkline-area" :d="getSparklineArea(savingsHistory)"></path>
<path class="sparkline" :d="getSparkline(savingsHistory)"></path>
</svg>
</div>
</div>
<!-- Cost Saved -->
<div class="bg-surface rounded-lg p-4 border border-border">
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">Cost Saved</div>
<div class="flex items-baseline gap-2">
<span class="text-3xl font-light tabular-nums text-accent" x-text="'$' + formatCost(stats.cost?.total_savings_usd || 0)"></span>
</div>
<div class="mt-3 text-xs text-gray-500">
vs $<span x-text="formatCost(stats.cost?.total_cost_usd || 0)"></span> spent
</div>
</div>
<!-- Headroom Overhead -->
<div class="bg-surface rounded-lg p-4 border border-border">
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">Headroom Overhead</div>
<div class="flex items-baseline gap-2">
<span class="text-3xl font-light tabular-nums" x-text="(stats.overhead?.average_ms || 0).toFixed(0) + 'ms'"></span>
</div>
<div class="mt-3 text-xs text-gray-500">
Avg <span x-text="((stats.latency?.average_ms || 0) / 1000).toFixed(1)"></span>s total response time
</div>
</div>
</div>
<!-- Main Grid -->
<div class="grid grid-cols-1 lg:grid-cols-3 gap-4 mb-6">
<!-- Token Usage -->
<div class="bg-surface rounded-lg p-4 border border-border">
<div class="text-sm font-medium mb-4 text-gray-300">Token Usage</div>
<div class="space-y-3">
<div class="flex justify-between items-center">
<span class="text-sm text-gray-400">Input Tokens</span>
<span class="font-mono text-sm" x-text="formatNumber(stats.tokens?.input || 0)"></span>
</div>
<div class="flex justify-between items-center">
<span class="text-sm text-gray-400">Output Tokens</span>
<span class="font-mono text-sm" x-text="formatNumber(stats.tokens?.output || 0)"></span>
</div>
<div class="border-t border-border my-2"></div>
<div class="flex justify-between items-center">
<span class="text-sm text-gray-400">Original Size</span>
<span class="font-mono text-sm" x-text="formatNumber(stats.tokens?.total_before_compression || 0)"></span>
</div>
<div class="flex justify-between items-center">
<span class="text-sm text-gray-400">After Compression</span>
<span class="font-mono text-sm" x-text="formatNumber(stats.tokens?.input || 0)"></span>
</div>
</div>
</div>
<!-- Providers -->
<div class="bg-surface rounded-lg p-4 border border-border">
<div class="text-sm font-medium mb-4 text-gray-300">Providers</div>
<div class="space-y-2">
<template x-for="(count, provider) in (stats.requests?.by_provider || {})" :key="provider">
<div class="flex justify-between items-center">
<span class="text-sm text-gray-400" x-text="provider"></span>
<div class="flex items-center gap-2">
<div class="w-24 h-1.5 bg-border rounded-full overflow-hidden">
<div class="h-full bg-accent rounded-full"
:style="'width: ' + getProviderPercent(count) + '%'"></div>
</div>
<span class="font-mono text-sm w-12 text-right" x-text="formatNumber(count)"></span>
</div>
</div>
</template>
<template x-if="Object.keys(stats.requests?.by_provider || {}).length === 0">
<div class="text-sm text-gray-500 italic">No requests yet</div>
</template>
</div>
</div>
<!-- Performance -->
<div class="bg-surface rounded-lg p-4 border border-border">
<div class="text-sm font-medium mb-4 text-gray-300">Performance</div>
<div class="space-y-3">
<div class="flex justify-between items-center">
<span class="text-sm text-gray-400">Headroom Overhead</span>
<span class="font-mono text-sm" x-text="(stats.overhead?.average_ms || 0).toFixed(0) + 'ms avg'"></span>
</div>
<div class="flex justify-between items-center">
<span class="text-sm text-gray-400">Overhead Range</span>
<span class="font-mono text-sm" x-text="(stats.overhead?.min_ms || 0).toFixed(0) + ' - ' + (stats.overhead?.max_ms || 0).toFixed(0) + 'ms'"></span>
</div>
<div class="flex justify-between items-center">
<span class="text-sm text-gray-400">Total Response Time</span>
<span class="font-mono text-sm" x-text="((stats.latency?.average_ms || 0) / 1000).toFixed(1) + 's avg'"></span>
</div>
<div class="flex justify-between items-center">
<span class="text-sm text-gray-400">Failed Requests</span>
<span class="font-mono text-sm" x-text="stats.requests?.failed || 0"></span>
</div>
</div>
</div>
</div>
<!-- Recent Requests Table -->
<div class="bg-surface rounded-lg border border-border overflow-hidden">
<div class="px-4 py-3 border-b border-border flex justify-between items-center">
<span class="text-sm font-medium text-gray-300">Recent Requests</span>
<span class="text-xs text-gray-500">Last 10</span>
</div>
<div class="overflow-x-auto">
<table class="w-full text-sm">
<thead>
<tr class="text-left text-xs text-gray-500 uppercase tracking-wide">
<th class="px-4 py-3 font-medium">Time</th>
<th class="px-4 py-3 font-medium">Model</th>
<th class="px-4 py-3 font-medium text-right">Input</th>
<th class="px-4 py-3 font-medium text-right">Output</th>
<th class="px-4 py-3 font-medium text-right">Saved</th>
<th class="px-4 py-3 font-medium text-right">Cost</th>
<th class="px-4 py-3 font-medium text-right">Latency</th>
</tr>
</thead>
<tbody class="divide-y divide-border">
<template x-for="req in (stats.recent_requests || [])" :key="req.request_id">
<tr class="hover:bg-border/30 transition-colors">
<td class="px-4 py-3 font-mono text-gray-400" x-text="formatTime(req.timestamp)"></td>
<td class="px-4 py-3">
<span class="px-2 py-0.5 bg-border rounded text-xs" x-text="truncateModel(req.model)"></span>
</td>
<td class="px-4 py-3 text-right font-mono" x-text="formatNumber(req.input_tokens_optimized)"></td>
<td class="px-4 py-3 text-right font-mono" x-text="formatNumber(req.output_tokens || 0)"></td>
<td class="px-4 py-3 text-right">
<span class="text-accent font-mono" x-text="req.savings_percent.toFixed(0) + '%'"></span>
</td>
<td class="px-4 py-3 text-right font-mono text-gray-400" x-text="'$' + (req.estimated_cost_usd || 0).toFixed(4)"></td>
<td class="px-4 py-3 text-right font-mono text-gray-400" x-text="(req.total_latency_ms || 0).toFixed(0) + 'ms'"></td>
</tr>
</template>
<template x-if="(stats.recent_requests || []).length === 0">
<tr>
<td colspan="7" class="px-4 py-8 text-center text-gray-500 italic">
No requests yet. Start using the proxy to see activity here.
</td>
</tr>
</template>
</tbody>
</table>
</div>
</div>
<!-- Budget Bar (if configured) -->
<template x-if="stats.cost?.budget_limit_usd">
<div class="mt-6 bg-surface rounded-lg p-4 border border-border">
<div class="flex justify-between items-center mb-2">
<span class="text-sm text-gray-400">Budget (<span x-text="stats.cost?.budget_period || 'daily'"></span>)</span>
<span class="font-mono text-sm">
$<span x-text="formatCost(stats.cost?.period_cost_usd || 0)"></span>
/ $<span x-text="formatCost(stats.cost?.budget_limit_usd || 0)"></span>
</span>
</div>
<div class="w-full h-2 bg-border rounded-full overflow-hidden">
<div class="h-full rounded-full transition-all duration-500"
:class="getBudgetPercent() > 90 ? 'bg-red-400' : getBudgetPercent() > 70 ? 'bg-amber-400' : 'bg-accent'"
:style="'width: ' + Math.min(getBudgetPercent(), 100) + '%'"></div>
</div>
</div>
</template>
</main>
<!-- Footer -->
<footer class="border-t border-border px-6 py-4 mt-8">
<div class="flex justify-between items-center text-xs text-gray-500">
<div>
Press <kbd class="px-1.5 py-0.5 bg-border rounded text-gray-400">R</kbd> to refresh
</div>
<div>
<a href="https://chopratejas.github.io/headroom/" target="_blank" class="hover:text-gray-300 transition-colors">Documentation</a>
</div>
</div>
</footer>
<script>
function dashboard() {
return {
stats: {},
healthy: true,
version: '0.3.0',
lastUpdate: 'never',
requestHistory: [],
savingsHistory: [],
pollInterval: null,
async init() {
await this.fetchStats();
this.pollInterval = setInterval(() => this.fetchStats(), 3000);
// Keyboard shortcuts
document.addEventListener('keydown', (e) => {
if (e.key === 'r' || e.key === 'R') {
this.fetchStats();
}
});
},
async fetchStats() {
try {
const [statsRes, healthRes] = await Promise.all([
fetch('/stats'),
fetch('/health')
]);
this.stats = await statsRes.json();
const health = await healthRes.json();
this.healthy = health.status === 'healthy';
this.version = health.version || '0.3.0';
// Update history for sparklines
this.requestHistory.push(this.stats.requests?.total || 0);
this.savingsHistory.push(this.stats.tokens?.saved || 0);
// Keep last 30 points
if (this.requestHistory.length > 30) this.requestHistory.shift();
if (this.savingsHistory.length > 30) this.savingsHistory.shift();
this.lastUpdate = new Date().toLocaleTimeString();
} catch (e) {
console.error('Failed to fetch stats:', e);
this.healthy = false;
}
},
formatNumber(n) {
if (n >= 1000000) return (n / 1000000).toFixed(1) + 'M';
if (n >= 1000) return (n / 1000).toFixed(1) + 'k';
return n.toString();
},
formatCost(n) {
return n.toFixed(2);
},
formatTime(ts) {
if (!ts) return '-';
const d = new Date(ts);
const now = new Date();
const diff = (now - d) / 1000;
if (diff < 60) return Math.floor(diff) + 's ago';
if (diff < 3600) return Math.floor(diff / 60) + 'm ago';
return d.toLocaleTimeString();
},
truncateModel(model) {
if (!model) return '-';
// Remove provider prefix and version suffix for display
return model.replace(/^(anthropic\.|openai\.|bedrock\/)/, '')
.replace(/-\d{8}$/, '')
.substring(0, 20);
},
getProviderPercent(count) {
const total = this.stats.requests?.total || 1;
return Math.min((count / total) * 100, 100);
},
getBudgetPercent() {
const limit = this.stats.cost?.budget_limit_usd || 1;
const used = this.stats.cost?.period_cost_usd || 0;
return (used / limit) * 100;
},
getSparkline(data) {
if (!data || data.length < 2) return '';
const min = Math.min(...data);
const max = Math.max(...data);
const range = max - min || 1;
const points = data.map((v, i) => {
const x = (i / (data.length - 1)) * 100;
const y = 32 - ((v - min) / range) * 28;
return `${x},${y}`;
});
return 'M' + points.join(' L');
},
getSparklineArea(data) {
if (!data || data.length < 2) return '';
const line = this.getSparkline(data);
if (!line) return '';
return line + ` L100,32 L0,32 Z`;
}
};
}
</script>
</body>
</html>

88
headroom/integrations/strands/__init__.py Normal file
View file

@ -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",
]

540
headroom/integrations/strands/hooks.py Normal file
View file

@ -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")

625
headroom/integrations/strands/model.py Normal file
View file

@ -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

166
headroom/integrations/strands/providers.py Normal file
View file

@ -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"

128
headroom/proxy/server.py
View file

@ -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),

4
pyproject.toml
View file

@ -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

1
tests/integrations/test_strands/__init__.py Normal file
View file

@ -0,0 +1 @@
"""Tests for Strands Agents SDK integration with Headroom."""

545
tests/integrations/test_strands/test_hooks.py Normal file
View file

@ -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

564
tests/integrations/test_strands/test_hooks_unit.py Normal file
View file

@ -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"}]

673
tests/integrations/test_strands/test_model.py Normal file
View file

@ -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

645
tests/integrations/test_strands/test_model_unit.py Normal file
View file

@ -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