chopratejas commited on
Commit
2550bb7
·
1 Parent(s): c9daa5d

feat: Add AWS Strands Agents SDK integration

Browse files

## Description

Add Headroom integration with AWS Strands Agents SDK, enabling automatic
context optimization and tool output compression for Strands-based agents.

Fixes #14

## Type of Change

- [x] New feature (non-breaking change that adds functionality)
- [x] Documentation update

## Changes Made

### Core Integration (`headroom/integrations/strands/`)

- **HeadroomHookProvider** - Implements Strands `HookProvider` interface for
automatic tool output compression via `AfterToolCallEvent`. Compresses
verbose tool outputs before they enter conversation context.

- **HeadroomStrandsModel** - Model wrapper that extends Strands `Model` base
class for message-level optimization. Implements all required abstract
methods: `stream()`, `get_config()`, `update_config()`, `structured_output()`.

- **Provider auto-detection** - Automatically detects appropriate Headroom
provider (Anthropic, OpenAI, Google) based on wrapped Strands model type.

- **`strands-agents` as optional dependency** - Install with
`pip install headroom-ai[strands]`

### Testing (`tests/integrations/test_strands/`)

- **Real integration tests (25 tests)** - Use actual AWS Bedrock API calls
with Claude 3 Haiku. Skip automatically when credentials unavailable.

- **Unit tests (57 tests)** - Mock-based tests for internal logic, edge cases,
and error handling. No credentials required.

### Demo (`examples/strands_bedrock_demo.py`)

- Interactive demo showcasing both integration patterns
- Visual before/after compression comparison with token savings
- 4 verbose tools (search, logs, database, metrics) demonstrating real savings
- Supports `--hook` and `--model` flags for individual demos

## Testing

All tests verified:

- [x] Unit tests pass (57 tests)
- [x] Integration tests pass (25 tests with real Bedrock API)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom/integrations/strands/`)
- [x] Formatting passes (`ruff format --check`)
- [x] Demo runs successfully with ~50% token savings

## Test Output

```
$ pytest tests/integrations/test_strands/ -v
=================== 82 passed in 90.09s ===================

$ ruff check headroom/integrations/strands/ --ignore E402
All checks passed!

$ mypy headroom/integrations/strands/ --ignore-missing-imports
Success: no issues found
```

## Demo Results

```
╭────────────────────────────────────────────────────────────╮
│ HeadroomHookProvider Results │
│────────────────────────────────────────────────────────────│
│ Tokens BEFORE compression: 51,961 │
│ Tokens AFTER compression: 25,658 │
│ Tokens SAVED: 26,303 (50.6%) │
╰────────────────────────────────────────────────────────────╯
```

examples/README.md CHANGED
@@ -86,6 +86,44 @@ export OPENAI_API_KEY='your-key'
86
  PYTHONPATH=. python -m examples.mcp_demo.run_agent_eval
87
  ```
88
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
89
  ## Running Examples
90
 
91
  All examples can be run from the repository root:
@@ -105,6 +143,7 @@ python examples/<example_name>.py
105
  | basic_usage | 50-70% | Simple tool output compression |
106
  | langchain_demo | 70-85% | Real agent with multiple tools |
107
  | mcp_demo | 60-80% | MCP tool outputs |
 
108
  | real_world_eval | 50-90% | Varies by scenario |
109
 
110
  ## Troubleshooting
@@ -131,3 +170,20 @@ Ensure your API keys are set:
131
  export OPENAI_API_KEY='sk-...'
132
  export ANTHROPIC_API_KEY='sk-ant-...'
133
  ```
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
86
  PYTHONPATH=. python -m examples.mcp_demo.run_agent_eval
87
  ```
88
 
89
+ ### strands_bedrock_demo.py
90
+
91
+ AWS Strands Agents + Bedrock integration demo. Showcases two Headroom integration patterns:
92
+
93
+ 1. **HeadroomHookProvider** - Compresses tool outputs in real-time
94
+ 2. **HeadroomStrandsModel** - Optimizes entire conversation context
95
+
96
+ ```bash
97
+ # Configure AWS credentials
98
+ export AWS_ACCESS_KEY_ID='your-access-key'
99
+ export AWS_SECRET_ACCESS_KEY='your-secret-key'
100
+ export AWS_DEFAULT_REGION='us-west-2' # Optional, defaults to us-west-2
101
+
102
+ # Or use AWS profile
103
+ export AWS_PROFILE='your-profile-name'
104
+
105
+ # Run the full demo (both integration patterns)
106
+ python examples/strands_bedrock_demo.py
107
+
108
+ # Run only the hook provider demo
109
+ python examples/strands_bedrock_demo.py --hook
110
+
111
+ # Run only the model wrapper demo
112
+ python examples/strands_bedrock_demo.py --model
113
+
114
+ # Specify a different AWS region
115
+ python examples/strands_bedrock_demo.py --region us-east-1
116
+ ```
117
+
118
+ The demo uses Claude 3 Haiku via Bedrock for cost efficiency. It creates agents with
119
+ 4 tools that return verbose JSON output (search results, logs, database records, metrics)
120
+ and displays compression statistics with visual comparisons.
121
+
122
+ **Requirements:**
123
+ - AWS account with Bedrock enabled
124
+ - Claude 3 Haiku model access in your region
125
+ - `pip install strands-agents headroom-ai[strands]`
126
+
127
  ## Running Examples
128
 
129
  All examples can be run from the repository root:
 
143
  | basic_usage | 50-70% | Simple tool output compression |
144
  | langchain_demo | 70-85% | Real agent with multiple tools |
145
  | mcp_demo | 60-80% | MCP tool outputs |
146
+ | strands_bedrock_demo | 60-85% | Strands + Bedrock with verbose tools |
147
  | real_world_eval | 50-90% | Varies by scenario |
148
 
149
  ## Troubleshooting
 
170
  export OPENAI_API_KEY='sk-...'
171
  export ANTHROPIC_API_KEY='sk-ant-...'
172
  ```
173
+
174
+ **AWS Credentials Errors (for Strands demo)**
175
+
176
+ Ensure AWS credentials are configured:
177
+
178
+ ```bash
179
+ # Option 1: Environment variables
180
+ export AWS_ACCESS_KEY_ID='your-access-key'
181
+ export AWS_SECRET_ACCESS_KEY='your-secret-key'
182
+
183
+ # Option 2: AWS profile
184
+ export AWS_PROFILE='your-profile-name'
185
+
186
+ # Option 3: AWS credentials file (~/.aws/credentials)
187
+ ```
188
+
189
+ Also ensure Bedrock and the Claude 3 Haiku model are enabled in your AWS account.
examples/strands_bedrock_demo.py ADDED
@@ -0,0 +1,1001 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Comprehensive Strands + Bedrock Demo for Headroom SDK.
3
+
4
+ This demo showcases two Headroom integration patterns for AWS Strands Agents:
5
+
6
+ 1. **HeadroomHookProvider** - Compresses tool outputs as they happen
7
+ - Intercepts tool results via Strands hooks
8
+ - Applies SmartCrusher compression to large JSON outputs
9
+ - Shows per-tool compression metrics
10
+
11
+ 2. **HeadroomStrandsModel** - Optimizes entire conversation context
12
+ - Wraps BedrockModel for automatic context optimization
13
+ - Applies message-level transforms before API calls
14
+ - Tracks cumulative savings across the session
15
+
16
+ Run with:
17
+ python examples/strands_bedrock_demo.py # Run both demos
18
+ python examples/strands_bedrock_demo.py --hook # Hook provider demo only
19
+ python examples/strands_bedrock_demo.py --model # Model wrapper demo only
20
+
21
+ Requirements:
22
+ - AWS credentials configured (AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY or AWS_PROFILE)
23
+ - pip install strands-agents headroom-ai[strands]
24
+ """
25
+
26
+ from __future__ import annotations
27
+
28
+ import argparse
29
+ import json
30
+ import os
31
+ import random
32
+ import sys
33
+ from datetime import datetime, timedelta
34
+ from typing import Any
35
+
36
+ # ============================================================================
37
+ # Check Dependencies
38
+ # ============================================================================
39
+
40
+
41
+ def check_dependencies() -> bool:
42
+ """Check if required dependencies are available."""
43
+ missing = []
44
+
45
+ # Check strands-agents
46
+ try:
47
+ from strands import Agent # noqa: F401
48
+ from strands.models import BedrockModel # noqa: F401
49
+ except ImportError:
50
+ missing.append("strands-agents")
51
+
52
+ # Check headroom
53
+ try:
54
+ from headroom.integrations.strands import ( # noqa: F401
55
+ HeadroomHookProvider,
56
+ HeadroomStrandsModel,
57
+ )
58
+ except ImportError:
59
+ missing.append("headroom-ai[strands]")
60
+
61
+ if missing:
62
+ print_box(
63
+ "Missing Dependencies",
64
+ [
65
+ "The following packages are required but not installed:",
66
+ "",
67
+ *[f" - {pkg}" for pkg in missing],
68
+ "",
69
+ "Install with:",
70
+ f" pip install {' '.join(missing)}",
71
+ ],
72
+ style="error",
73
+ )
74
+ return False
75
+
76
+ return True
77
+
78
+
79
+ def check_aws_credentials() -> bool:
80
+ """Check if AWS credentials are available."""
81
+ has_env_keys = os.environ.get("AWS_ACCESS_KEY_ID") and os.environ.get("AWS_SECRET_ACCESS_KEY")
82
+ has_profile = os.environ.get("AWS_PROFILE")
83
+ has_creds_file = os.path.exists(os.path.expanduser("~/.aws/credentials"))
84
+
85
+ if not (has_env_keys or has_profile or has_creds_file):
86
+ print_box(
87
+ "AWS Credentials Not Found",
88
+ [
89
+ "This demo requires AWS credentials to access Bedrock.",
90
+ "",
91
+ "Configure credentials using one of these methods:",
92
+ "",
93
+ "1. Environment variables:",
94
+ " export AWS_ACCESS_KEY_ID='your-access-key'",
95
+ " export AWS_SECRET_ACCESS_KEY='your-secret-key'",
96
+ " export AWS_DEFAULT_REGION='us-west-2'",
97
+ "",
98
+ "2. AWS Profile:",
99
+ " export AWS_PROFILE='your-profile-name'",
100
+ "",
101
+ "3. AWS credentials file:",
102
+ " ~/.aws/credentials",
103
+ ],
104
+ style="error",
105
+ )
106
+ return False
107
+
108
+ return True
109
+
110
+
111
+ # ============================================================================
112
+ # Pretty Printing Utilities
113
+ # ============================================================================
114
+
115
+
116
+ def print_box(title: str, lines: list[str], style: str = "normal", width: int = 76) -> None:
117
+ """Print a box with title and content using box drawing characters."""
118
+ if style == "error":
119
+ top_left, top_right = "\u2554", "\u2557" # Double line
120
+ bot_left, bot_right = "\u255a", "\u255d"
121
+ horiz, vert = "\u2550", "\u2551"
122
+ elif style == "success":
123
+ top_left, top_right = "\u256d", "\u256e" # Rounded
124
+ bot_left, bot_right = "\u2570", "\u256f"
125
+ horiz, vert = "\u2500", "\u2502"
126
+ else:
127
+ top_left, top_right = "\u250c", "\u2510" # Normal single
128
+ bot_left, bot_right = "\u2514", "\u2518"
129
+ horiz, vert = "\u2500", "\u2502"
130
+
131
+ print()
132
+ print(f"{top_left}{horiz * (width - 2)}{top_right}")
133
+
134
+ # Title
135
+ title_padding = (width - 4 - len(title)) // 2
136
+ print(
137
+ f"{vert} {' ' * title_padding}{title}{' ' * (width - 4 - title_padding - len(title))} {vert}"
138
+ )
139
+ print(f"{vert}{horiz * (width - 2)}{vert}")
140
+
141
+ # Content lines
142
+ for line in lines:
143
+ # Handle lines longer than width
144
+ if len(line) > width - 4:
145
+ line = line[: width - 7] + "..."
146
+ padding = width - 4 - len(line)
147
+ print(f"{vert} {line}{' ' * padding} {vert}")
148
+
149
+ print(f"{bot_left}{horiz * (width - 2)}{bot_right}")
150
+ print()
151
+
152
+
153
+ def print_metrics_table(
154
+ metrics: list[dict[str, Any]],
155
+ headers: list[str],
156
+ keys: list[str],
157
+ title: str = "Metrics",
158
+ ) -> None:
159
+ """Print metrics in a formatted table."""
160
+ # Calculate column widths
161
+ col_widths = []
162
+ for i, header in enumerate(headers):
163
+ max_width = len(header)
164
+ for m in metrics:
165
+ val = m.get(keys[i], "")
166
+ max_width = max(max_width, len(str(val)))
167
+ col_widths.append(min(max_width + 2, 25))
168
+
169
+ total_width = sum(col_widths) + len(col_widths) + 1
170
+
171
+ print(f"\n {title}")
172
+ print(" " + "\u2500" * (total_width - 2))
173
+
174
+ # Header row
175
+ header_row = "\u2502"
176
+ for i, header in enumerate(headers):
177
+ header_row += f" {header:<{col_widths[i] - 2}} \u2502"
178
+ print(" " + header_row)
179
+ print(" " + "\u2502" + "\u2500" * (total_width - 2) + "\u2502")
180
+
181
+ # Data rows
182
+ for m in metrics:
183
+ row = "\u2502"
184
+ for i, key in enumerate(keys):
185
+ val = str(m.get(key, ""))
186
+ if len(val) > col_widths[i] - 2:
187
+ val = val[: col_widths[i] - 5] + "..."
188
+ row += f" {val:<{col_widths[i] - 2}} \u2502"
189
+ print(" " + row)
190
+
191
+ print(" " + "\u2500" * total_width)
192
+
193
+
194
+ def print_comparison(before: int, after: int, label: str = "Tokens") -> None:
195
+ """Print a before/after comparison with savings."""
196
+ saved = before - after
197
+ pct = (saved / before * 100) if before > 0 else 0
198
+
199
+ bar_width = 40
200
+ before_bar = int((before / max(before, 1)) * bar_width)
201
+ after_bar = int((after / max(before, 1)) * bar_width)
202
+
203
+ print(f"\n {label} Comparison:")
204
+ print(f" BEFORE: {before:>8,} \u2502{'=' * before_bar}")
205
+ print(f" AFTER: {after:>8,} \u2502{'=' * after_bar}")
206
+ print(f" SAVED: {saved:>8,} ({pct:.1f}%)")
207
+
208
+
209
+ # ============================================================================
210
+ # Mock Tools - Generate Verbose Output
211
+ # ============================================================================
212
+
213
+
214
+ def search_documentation(query: str, limit: int = 25) -> str:
215
+ """Search documentation for matching articles.
216
+
217
+ Returns search results with titles, snippets, URLs, and metadata.
218
+ Simulates a real documentation search API returning verbose results.
219
+ """
220
+ results = []
221
+ categories = [
222
+ "getting-started",
223
+ "api-reference",
224
+ "tutorials",
225
+ "troubleshooting",
226
+ "best-practices",
227
+ ]
228
+ sources = ["internal-docs", "confluence", "notion", "github-wiki", "readme"]
229
+
230
+ for i in range(limit):
231
+ result = {
232
+ "id": f"doc-{random.randint(10000, 99999)}",
233
+ "title": f"{query.title()} Guide - Part {i + 1}",
234
+ "snippet": f"This comprehensive guide covers {query} implementation. "
235
+ f"Learn how to configure, deploy, and maintain {query} in production. "
236
+ f"Includes examples, best practices, and troubleshooting tips for {query}.",
237
+ "url": f"https://docs.example.com/{query.replace(' ', '-')}/section-{i + 1}",
238
+ "category": random.choice(categories),
239
+ "source": random.choice(sources),
240
+ "relevance_score": round(random.uniform(0.5, 1.0), 3),
241
+ "last_updated": (datetime.now() - timedelta(days=random.randint(1, 180))).isoformat(),
242
+ "author": f"Author {random.randint(1, 20)}",
243
+ "word_count": random.randint(500, 5000),
244
+ "views": random.randint(100, 10000),
245
+ "helpful_votes": random.randint(10, 500),
246
+ "tags": random.sample(
247
+ ["aws", "python", "deployment", "security", "performance", "monitoring"],
248
+ k=random.randint(2, 4),
249
+ ),
250
+ }
251
+ results.append(result)
252
+
253
+ results.sort(key=lambda x: x["relevance_score"], reverse=True)
254
+
255
+ return json.dumps(
256
+ {
257
+ "query": query,
258
+ "total_results": limit * 5, # Simulate more results available
259
+ "page": 1,
260
+ "per_page": limit,
261
+ "results": results,
262
+ },
263
+ indent=2,
264
+ )
265
+
266
+
267
+ def get_server_logs(server: str, lines: int = 100) -> str:
268
+ """Fetch server logs for analysis.
269
+
270
+ Returns JSON log entries with timestamps, levels, messages, and context.
271
+ Simulates verbose application logs with mostly INFO entries and some errors.
272
+ """
273
+ entries = []
274
+ levels = ["DEBUG", "INFO", "INFO", "INFO", "INFO", "WARN", "ERROR"]
275
+ services = ["api-gateway", "auth-service", "data-processor", "cache-layer", "message-queue"]
276
+
277
+ for _i in range(lines):
278
+ timestamp = datetime.now() - timedelta(minutes=random.randint(1, 1440))
279
+ level = random.choice(levels)
280
+
281
+ if level == "ERROR":
282
+ message = random.choice(
283
+ [
284
+ f"Connection timeout to {server}-db after 30000ms",
285
+ "Failed to authenticate request: invalid JWT signature",
286
+ "Rate limit exceeded for client IP 10.0.0.42",
287
+ "Database query failed: connection pool exhausted",
288
+ f"Service {server} health check failed: connection refused",
289
+ ]
290
+ )
291
+ elif level == "WARN":
292
+ message = random.choice(
293
+ [
294
+ f"Slow query detected on {server}: execution time 2.5s",
295
+ "Memory usage at 85% - consider scaling",
296
+ "Retry attempt 2/3 for downstream service call",
297
+ "Certificate expires in 7 days - renewal required",
298
+ ]
299
+ )
300
+ else:
301
+ message = f"Request processed successfully - endpoint=/api/v1/{server}/data"
302
+
303
+ entry = {
304
+ "timestamp": timestamp.isoformat(),
305
+ "level": level,
306
+ "server": server,
307
+ "service": random.choice(services),
308
+ "message": message,
309
+ "trace_id": f"trace-{random.randint(100000, 999999):06x}",
310
+ "span_id": f"span-{random.randint(1000, 9999):04x}",
311
+ "request_id": f"req-{random.randint(10000000, 99999999)}",
312
+ "client_ip": f"10.0.{random.randint(0, 255)}.{random.randint(1, 254)}",
313
+ "user_agent": random.choice(
314
+ [
315
+ "Mozilla/5.0 (compatible; MonitorBot/1.0)",
316
+ "python-requests/2.31.0",
317
+ "curl/8.1.2",
318
+ "PostmanRuntime/7.32.0",
319
+ ]
320
+ ),
321
+ "response_time_ms": random.randint(5, 2000),
322
+ "status_code": 200
323
+ if level in ["DEBUG", "INFO"]
324
+ else random.choice([400, 500, 502, 503]),
325
+ "metadata": {
326
+ "pod": f"{server}-{random.randint(1, 5)}-abc123",
327
+ "node": f"ip-10-0-{random.randint(0, 255)}-{random.randint(1, 254)}.ec2.internal",
328
+ "region": random.choice(["us-west-2", "us-east-1", "eu-west-1"]),
329
+ "version": f"v1.{random.randint(0, 9)}.{random.randint(0, 20)}",
330
+ },
331
+ }
332
+ entries.append(entry)
333
+
334
+ entries.sort(key=lambda x: x["timestamp"], reverse=True)
335
+
336
+ return json.dumps(
337
+ {
338
+ "server": server,
339
+ "log_count": lines,
340
+ "time_range": {
341
+ "start": entries[-1]["timestamp"] if entries else None,
342
+ "end": entries[0]["timestamp"] if entries else None,
343
+ },
344
+ "entries": entries,
345
+ },
346
+ indent=2,
347
+ )
348
+
349
+
350
+ def query_database(sql: str, limit: int = 50) -> str:
351
+ """Execute a database query and return results.
352
+
353
+ Returns rows of data as if from a real database query.
354
+ Simulates customer/order/transaction data.
355
+ """
356
+ # Parse table name from SQL (simple simulation)
357
+ table = "records"
358
+ for word in sql.lower().split():
359
+ if word in ["users", "orders", "transactions", "customers", "products", "events"]:
360
+ table = word
361
+ break
362
+
363
+ rows = []
364
+ statuses = ["active", "pending", "completed", "cancelled", "refunded"]
365
+
366
+ for i in range(limit):
367
+ if table == "users":
368
+ row = {
369
+ "user_id": f"usr-{random.randint(100000, 999999)}",
370
+ "email": f"user{i}@example.com",
371
+ "name": f"Customer {i}",
372
+ "status": random.choice(["active", "inactive", "suspended"]),
373
+ "created_at": (datetime.now() - timedelta(days=random.randint(1, 365))).isoformat(),
374
+ "last_login": (
375
+ datetime.now() - timedelta(hours=random.randint(1, 720))
376
+ ).isoformat(),
377
+ "plan": random.choice(["free", "basic", "pro", "enterprise"]),
378
+ "country": random.choice(["US", "UK", "DE", "FR", "JP", "AU"]),
379
+ }
380
+ elif table == "orders":
381
+ row = {
382
+ "order_id": f"ord-{random.randint(100000, 999999)}",
383
+ "customer_id": f"usr-{random.randint(100000, 999999)}",
384
+ "total": round(random.uniform(10, 1000), 2),
385
+ "currency": random.choice(["USD", "EUR", "GBP"]),
386
+ "status": random.choice(statuses),
387
+ "items_count": random.randint(1, 10),
388
+ "created_at": (datetime.now() - timedelta(days=random.randint(1, 90))).isoformat(),
389
+ "shipped_at": (datetime.now() - timedelta(days=random.randint(0, 30))).isoformat()
390
+ if random.random() > 0.3
391
+ else None,
392
+ }
393
+ else:
394
+ row = {
395
+ "id": i + 1,
396
+ "record_type": table,
397
+ "value": random.randint(100, 10000),
398
+ "status": random.choice(statuses),
399
+ "created_at": (datetime.now() - timedelta(days=random.randint(1, 365))).isoformat(),
400
+ "metadata": {
401
+ "source": random.choice(["web", "api", "import", "sync"]),
402
+ "version": f"v{random.randint(1, 5)}",
403
+ },
404
+ }
405
+ rows.append(row)
406
+
407
+ return json.dumps(
408
+ {
409
+ "query": sql,
410
+ "table": table,
411
+ "row_count": limit,
412
+ "total_available": limit * 10,
413
+ "execution_time_ms": random.randint(10, 500),
414
+ "rows": rows,
415
+ },
416
+ indent=2,
417
+ )
418
+
419
+
420
+ def get_system_metrics(timerange: str = "1h", service: str = "all") -> str:
421
+ """Get system metrics for monitoring.
422
+
423
+ Returns time-series data points for CPU, memory, latency, and error rates.
424
+ Simulates Prometheus/CloudWatch style metrics.
425
+ """
426
+ # Parse timerange to determine number of points
427
+ points = {"5m": 10, "15m": 30, "1h": 60, "6h": 72, "24h": 144}.get(timerange, 60)
428
+
429
+ data_points = []
430
+ services_list = ["api", "worker", "cache", "database"] if service == "all" else [service]
431
+
432
+ for svc in services_list:
433
+ for i in range(points):
434
+ timestamp = datetime.now() - timedelta(minutes=i * (60 // min(points, 60)))
435
+
436
+ # Inject some anomalies
437
+ is_anomaly = random.random() < 0.05
438
+
439
+ point = {
440
+ "timestamp": timestamp.isoformat(),
441
+ "service": svc,
442
+ "metrics": {
443
+ "cpu_percent": round(
444
+ random.uniform(70, 95) if is_anomaly else random.uniform(20, 45), 2
445
+ ),
446
+ "memory_percent": round(
447
+ random.uniform(80, 95) if is_anomaly else random.uniform(40, 65), 2
448
+ ),
449
+ "memory_mb": random.randint(2000, 4000)
450
+ if is_anomaly
451
+ else random.randint(500, 1500),
452
+ "latency_p50_ms": random.randint(100, 500)
453
+ if is_anomaly
454
+ else random.randint(10, 50),
455
+ "latency_p95_ms": random.randint(500, 2000)
456
+ if is_anomaly
457
+ else random.randint(50, 150),
458
+ "latency_p99_ms": random.randint(1000, 5000)
459
+ if is_anomaly
460
+ else random.randint(100, 300),
461
+ "request_rate_per_sec": random.randint(500, 2000)
462
+ if is_anomaly
463
+ else random.randint(50, 200),
464
+ "error_rate_percent": round(
465
+ random.uniform(5, 15) if is_anomaly else random.uniform(0, 1), 3
466
+ ),
467
+ "active_connections": random.randint(200, 500)
468
+ if is_anomaly
469
+ else random.randint(20, 80),
470
+ },
471
+ "health": "degraded" if is_anomaly else "healthy",
472
+ "region": random.choice(["us-west-2", "us-east-1", "eu-west-1"]),
473
+ }
474
+ data_points.append(point)
475
+
476
+ # Calculate summary statistics
477
+ all_cpu = [p["metrics"]["cpu_percent"] for p in data_points]
478
+ all_mem = [p["metrics"]["memory_percent"] for p in data_points]
479
+ all_latency = [p["metrics"]["latency_p50_ms"] for p in data_points]
480
+
481
+ return json.dumps(
482
+ {
483
+ "timerange": timerange,
484
+ "service": service,
485
+ "data_points_count": len(data_points),
486
+ "summary": {
487
+ "cpu": {
488
+ "min": min(all_cpu),
489
+ "max": max(all_cpu),
490
+ "avg": sum(all_cpu) / len(all_cpu),
491
+ },
492
+ "memory": {
493
+ "min": min(all_mem),
494
+ "max": max(all_mem),
495
+ "avg": sum(all_mem) / len(all_mem),
496
+ },
497
+ "latency_p50": {
498
+ "min": min(all_latency),
499
+ "max": max(all_latency),
500
+ "avg": sum(all_latency) / len(all_latency),
501
+ },
502
+ },
503
+ "data_points": data_points,
504
+ },
505
+ indent=2,
506
+ )
507
+
508
+
509
+ # ============================================================================
510
+ # Demo 1: HeadroomHookProvider
511
+ # ============================================================================
512
+
513
+
514
+ def run_hook_provider_demo(region: str = "us-west-2") -> dict[str, Any]:
515
+ """Demonstrate HeadroomHookProvider for tool output compression.
516
+
517
+ Returns metrics from the demo run.
518
+ """
519
+ from strands import Agent, tool
520
+ from strands.models import BedrockModel
521
+
522
+ from headroom.integrations.strands import HeadroomHookProvider
523
+
524
+ print_box(
525
+ "Demo 1: HeadroomHookProvider",
526
+ [
527
+ "The HeadroomHookProvider intercepts tool outputs and compresses",
528
+ "them BEFORE they're added to the conversation context.",
529
+ "",
530
+ "This reduces token usage for subsequent LLM calls by eliminating",
531
+ "redundant data from verbose tool outputs.",
532
+ "",
533
+ "Using: Claude 3 Haiku (anthropic.claude-3-haiku-20240307-v1:0)",
534
+ ],
535
+ )
536
+
537
+ # Define tools with @tool decorator
538
+ @tool
539
+ def search_docs_tool(query: str) -> str:
540
+ """Search documentation for articles matching the query.
541
+
542
+ Args:
543
+ query: The search query to find relevant documentation
544
+
545
+ Returns:
546
+ JSON array of search results with titles, snippets, and URLs
547
+ """
548
+ return search_documentation(query, limit=25)
549
+
550
+ @tool
551
+ def get_logs_tool(server: str, lines: int = 100) -> str:
552
+ """Fetch server logs for analysis and troubleshooting.
553
+
554
+ Args:
555
+ server: Name of the server to fetch logs from
556
+ lines: Number of log lines to retrieve (default: 100)
557
+
558
+ Returns:
559
+ JSON array of log entries with timestamps and messages
560
+ """
561
+ return get_server_logs(server, lines=lines)
562
+
563
+ @tool
564
+ def query_db_tool(sql: str) -> str:
565
+ """Execute a database query and return results.
566
+
567
+ Args:
568
+ sql: SQL query to execute (e.g., SELECT * FROM users)
569
+
570
+ Returns:
571
+ JSON array of database rows
572
+ """
573
+ return query_database(sql, limit=50)
574
+
575
+ @tool
576
+ def get_metrics_tool(timerange: str = "1h") -> str:
577
+ """Get system metrics for the specified time range.
578
+
579
+ Args:
580
+ timerange: Time range for metrics (5m, 15m, 1h, 6h, 24h)
581
+
582
+ Returns:
583
+ JSON object with time-series metrics data
584
+ """
585
+ return get_system_metrics(timerange)
586
+
587
+ # Create BedrockModel
588
+ model = BedrockModel(
589
+ model_id="anthropic.claude-3-haiku-20240307-v1:0",
590
+ region_name=region,
591
+ temperature=0.1,
592
+ )
593
+
594
+ # Create HeadroomHookProvider
595
+ hook_provider = HeadroomHookProvider(
596
+ compress_tool_outputs=True,
597
+ min_tokens_to_compress=100, # Compress outputs with 100+ tokens
598
+ preserve_errors=True,
599
+ )
600
+
601
+ # Create agent with hook
602
+ agent = Agent(
603
+ model=model,
604
+ tools=[search_docs_tool, get_logs_tool, query_db_tool, get_metrics_tool],
605
+ hooks=[hook_provider],
606
+ )
607
+
608
+ print("\n Running agent queries that trigger tools with verbose output...")
609
+ print(" " + "-" * 60)
610
+
611
+ # Query 1: Search documentation
612
+ print("\n Query 1: Searching documentation...")
613
+ result1 = agent(
614
+ "Search the documentation for 'authentication setup' and summarize "
615
+ "the top 3 most relevant articles you find."
616
+ )
617
+ print(f" Response: {str(result1)[:200]}...")
618
+
619
+ # Query 2: Get server logs
620
+ print("\n Query 2: Fetching server logs...")
621
+ result2 = agent(
622
+ "Get the logs from server 'api-gateway' (100 lines) and tell me "
623
+ "how many ERROR and WARN level entries there are."
624
+ )
625
+ print(f" Response: {str(result2)[:200]}...")
626
+
627
+ # Query 3: Query database
628
+ print("\n Query 3: Running database query...")
629
+ result3 = agent(
630
+ "Query the orders table and tell me how many orders have status 'completed' "
631
+ "and what the average order total is."
632
+ )
633
+ print(f" Response: {str(result3)[:200]}...")
634
+
635
+ # Query 4: Get metrics
636
+ print("\n Query 4: Fetching system metrics...")
637
+ result4 = agent(
638
+ "Get the system metrics for the last hour and identify if there are "
639
+ "any services with high CPU usage (>70%) or memory issues."
640
+ )
641
+ print(f" Response: {str(result4)[:200]}...")
642
+
643
+ # Get metrics
644
+ metrics = hook_provider.get_savings_summary()
645
+
646
+ # Display results
647
+ print_box(
648
+ "HeadroomHookProvider Results",
649
+ [
650
+ f"Tool calls processed: {metrics['total_requests']}",
651
+ f"Compressions applied: {metrics['compressed_requests']}",
652
+ "",
653
+ f"Tokens BEFORE compression: {metrics['total_tokens_before']:,}",
654
+ f"Tokens AFTER compression: {metrics['total_tokens_after']:,}",
655
+ f"Tokens SAVED: {metrics['total_tokens_saved']:,}",
656
+ "",
657
+ f"Average savings: {metrics['average_savings_percent']:.1f}%",
658
+ ],
659
+ style="success",
660
+ )
661
+
662
+ # Show per-tool breakdown
663
+ if hook_provider.metrics_history:
664
+ tool_metrics = []
665
+ for m in hook_provider.metrics_history:
666
+ tool_metrics.append(
667
+ {
668
+ "tool": m.tool_name[:20],
669
+ "before": f"{m.tokens_before:,}",
670
+ "after": f"{m.tokens_after:,}",
671
+ "saved": f"{m.tokens_saved:,}",
672
+ "pct": f"{m.savings_percent:.1f}%",
673
+ }
674
+ )
675
+
676
+ print_metrics_table(
677
+ tool_metrics,
678
+ headers=["Tool", "Before", "After", "Saved", "%"],
679
+ keys=["tool", "before", "after", "saved", "pct"],
680
+ title="Per-Tool Compression Breakdown",
681
+ )
682
+
683
+ print_comparison(
684
+ metrics["total_tokens_before"],
685
+ metrics["total_tokens_after"],
686
+ "Tool Output Tokens",
687
+ )
688
+
689
+ return metrics
690
+
691
+
692
+ # ============================================================================
693
+ # Demo 2: HeadroomStrandsModel
694
+ # ============================================================================
695
+
696
+
697
+ def run_model_wrapper_demo(region: str = "us-west-2") -> dict[str, Any]:
698
+ """Demonstrate HeadroomStrandsModel for conversation optimization.
699
+
700
+ Returns metrics from the demo run.
701
+ """
702
+ from strands import Agent, tool
703
+ from strands.models import BedrockModel
704
+
705
+ from headroom import HeadroomConfig
706
+ from headroom.integrations.strands import HeadroomStrandsModel
707
+
708
+ print_box(
709
+ "Demo 2: HeadroomStrandsModel",
710
+ [
711
+ "HeadroomStrandsModel wraps the Bedrock model to optimize the",
712
+ "ENTIRE conversation context before each API call.",
713
+ "",
714
+ "As conversations grow with tool outputs and history, the",
715
+ "model wrapper applies transforms to reduce context size.",
716
+ "",
717
+ "Using: Claude 3 Haiku wrapped with HeadroomStrandsModel",
718
+ ],
719
+ )
720
+
721
+ # Define tools
722
+ @tool
723
+ def verbose_search(query: str) -> str:
724
+ """Search for information with verbose results.
725
+
726
+ Args:
727
+ query: Search query
728
+
729
+ Returns:
730
+ Detailed search results
731
+ """
732
+ return search_documentation(query, limit=30)
733
+
734
+ @tool
735
+ def verbose_logs(server: str) -> str:
736
+ """Get verbose server logs.
737
+
738
+ Args:
739
+ server: Server name
740
+
741
+ Returns:
742
+ Detailed log entries
743
+ """
744
+ return get_server_logs(server, lines=150)
745
+
746
+ @tool
747
+ def verbose_metrics(timerange: str = "1h") -> str:
748
+ """Get verbose metrics data.
749
+
750
+ Args:
751
+ timerange: Time range
752
+
753
+ Returns:
754
+ Detailed metrics
755
+ """
756
+ return get_system_metrics(timerange)
757
+
758
+ @tool
759
+ def verbose_database(table: str) -> str:
760
+ """Query database with verbose results.
761
+
762
+ Args:
763
+ table: Table name to query
764
+
765
+ Returns:
766
+ Database records
767
+ """
768
+ return query_database(f"SELECT * FROM {table}", limit=60)
769
+
770
+ # Create base Bedrock model
771
+ base_model = BedrockModel(
772
+ model_id="anthropic.claude-3-haiku-20240307-v1:0",
773
+ region_name=region,
774
+ temperature=0.1,
775
+ )
776
+
777
+ # Configure Headroom
778
+ config = HeadroomConfig()
779
+ config.smart_crusher.enabled = True
780
+ config.smart_crusher.min_tokens_to_crush = 100
781
+ config.smart_crusher.max_items_after_crush = 20
782
+
783
+ # Wrap with HeadroomStrandsModel
784
+ optimized_model = HeadroomStrandsModel(
785
+ wrapped_model=base_model,
786
+ config=config,
787
+ auto_detect_provider=True,
788
+ )
789
+
790
+ # Create agent
791
+ agent = Agent(
792
+ model=optimized_model,
793
+ tools=[verbose_search, verbose_logs, verbose_metrics, verbose_database],
794
+ )
795
+
796
+ print("\n Building up a multi-turn conversation with verbose tool outputs...")
797
+ print(" " + "-" * 60)
798
+
799
+ # Simulate a multi-turn conversation
800
+ turns = [
801
+ ("Turn 1", "Search for documentation about 'kubernetes deployment' and give me a summary."),
802
+ ("Turn 2", "Now get the logs from the 'worker-service' server and identify any errors."),
803
+ ("Turn 3", "Query the orders database and tell me the distribution of order statuses."),
804
+ ("Turn 4", "Get the system metrics for the last hour and highlight any anomalies."),
805
+ ("Turn 5", "Based on everything you've found, what's the overall system health status?"),
806
+ ]
807
+
808
+ for turn_name, query in turns:
809
+ print(f"\n {turn_name}: {query[:60]}...")
810
+ result = agent(query)
811
+ print(f" Response: {str(result)[:150]}...")
812
+
813
+ # Get metrics
814
+ metrics = optimized_model.get_savings_summary()
815
+
816
+ # Display results
817
+ print_box(
818
+ "HeadroomStrandsModel Results",
819
+ [
820
+ f"API calls made: {metrics['total_requests']}",
821
+ "",
822
+ f"Total tokens BEFORE opt: {metrics['total_tokens_before']:,}",
823
+ f"Total tokens AFTER opt: {metrics['total_tokens_after']:,}",
824
+ f"Total tokens SAVED: {metrics['total_tokens_saved']:,}",
825
+ "",
826
+ f"Average savings per call: {metrics['average_savings_percent']:.1f}%",
827
+ ],
828
+ style="success",
829
+ )
830
+
831
+ # Show per-request breakdown
832
+ if optimized_model.metrics_history:
833
+ request_metrics = []
834
+ for i, m in enumerate(optimized_model.metrics_history):
835
+ request_metrics.append(
836
+ {
837
+ "request": f"Request {i + 1}",
838
+ "before": f"{m.tokens_before:,}",
839
+ "after": f"{m.tokens_after:,}",
840
+ "saved": f"{m.tokens_saved:,}",
841
+ "pct": f"{m.savings_percent:.1f}%",
842
+ }
843
+ )
844
+
845
+ print_metrics_table(
846
+ request_metrics,
847
+ headers=["Request", "Before", "After", "Saved", "%"],
848
+ keys=["request", "before", "after", "saved", "pct"],
849
+ title="Per-Request Optimization Breakdown",
850
+ )
851
+
852
+ print_comparison(
853
+ metrics["total_tokens_before"],
854
+ metrics["total_tokens_after"],
855
+ "Conversation Tokens",
856
+ )
857
+
858
+ return metrics
859
+
860
+
861
+ # ============================================================================
862
+ # Main
863
+ # ============================================================================
864
+
865
+
866
+ def main() -> int:
867
+ """Run the Strands Bedrock demo."""
868
+ parser = argparse.ArgumentParser(
869
+ description="Headroom + Strands Bedrock Demo",
870
+ formatter_class=argparse.RawDescriptionHelpFormatter,
871
+ epilog="""
872
+ Examples:
873
+ python examples/strands_bedrock_demo.py # Run both demos
874
+ python examples/strands_bedrock_demo.py --hook # Hook provider only
875
+ python examples/strands_bedrock_demo.py --model # Model wrapper only
876
+
877
+ Environment Variables:
878
+ AWS_ACCESS_KEY_ID AWS access key
879
+ AWS_SECRET_ACCESS_KEY AWS secret key
880
+ AWS_DEFAULT_REGION AWS region (default: us-west-2)
881
+ AWS_PROFILE AWS profile name (alternative to keys)
882
+ """,
883
+ )
884
+ parser.add_argument(
885
+ "--hook",
886
+ action="store_true",
887
+ help="Run only the HeadroomHookProvider demo",
888
+ )
889
+ parser.add_argument(
890
+ "--model",
891
+ action="store_true",
892
+ help="Run only the HeadroomStrandsModel demo",
893
+ )
894
+ parser.add_argument(
895
+ "--region",
896
+ default=os.environ.get("AWS_DEFAULT_REGION", "us-west-2"),
897
+ help="AWS region for Bedrock (default: us-west-2)",
898
+ )
899
+
900
+ args = parser.parse_args()
901
+
902
+ # If neither flag is set, run both
903
+ run_hook = args.hook or (not args.hook and not args.model)
904
+ run_model = args.model or (not args.hook and not args.model)
905
+
906
+ # Print header
907
+ print_box(
908
+ "Headroom + Strands Bedrock Demo",
909
+ [
910
+ "This demo showcases Headroom's integration with AWS Strands Agents.",
911
+ "",
912
+ "Headroom provides two integration patterns:",
913
+ " 1. HeadroomHookProvider - Compress tool outputs in real-time",
914
+ " 2. HeadroomStrandsModel - Optimize entire conversation context",
915
+ "",
916
+ f"Region: {args.region}",
917
+ "Model: Claude 3 Haiku (fast and cost-effective for demos)",
918
+ ],
919
+ )
920
+
921
+ # Check dependencies
922
+ if not check_dependencies():
923
+ return 1
924
+
925
+ # Check AWS credentials
926
+ if not check_aws_credentials():
927
+ return 1
928
+
929
+ print("\n All checks passed. Starting demos...\n")
930
+
931
+ all_metrics = {}
932
+
933
+ try:
934
+ # Run hook provider demo
935
+ if run_hook:
936
+ hook_metrics = run_hook_provider_demo(region=args.region)
937
+ all_metrics["hook_provider"] = hook_metrics
938
+
939
+ # Run model wrapper demo
940
+ if run_model:
941
+ model_metrics = run_model_wrapper_demo(region=args.region)
942
+ all_metrics["model_wrapper"] = model_metrics
943
+
944
+ # Print final summary
945
+ if run_hook and run_model:
946
+ total_before = all_metrics.get("hook_provider", {}).get(
947
+ "total_tokens_before", 0
948
+ ) + all_metrics.get("model_wrapper", {}).get("total_tokens_before", 0)
949
+ total_after = all_metrics.get("hook_provider", {}).get(
950
+ "total_tokens_after", 0
951
+ ) + all_metrics.get("model_wrapper", {}).get("total_tokens_after", 0)
952
+ total_saved = total_before - total_after
953
+ total_pct = (total_saved / total_before * 100) if total_before > 0 else 0
954
+
955
+ # Estimate cost savings (Claude 3 Haiku pricing)
956
+ # Input: $0.25 / 1M tokens, Output: $1.25 / 1M tokens
957
+ cost_per_token = 0.25 / 1_000_000
958
+ cost_saved = total_saved * cost_per_token
959
+
960
+ print_box(
961
+ "Session Summary",
962
+ [
963
+ "Combined metrics from both demos:",
964
+ "",
965
+ f"Total tokens processed: {total_before:,}",
966
+ f"Total tokens after opt: {total_after:,}",
967
+ f"Total tokens saved: {total_saved:,} ({total_pct:.1f}%)",
968
+ "",
969
+ f"Estimated cost savings: ${cost_saved:.6f}",
970
+ "(At scale, these savings compound significantly!)",
971
+ "",
972
+ "Integration patterns demonstrated:",
973
+ " [x] HeadroomHookProvider - Real-time tool output compression",
974
+ " [x] HeadroomStrandsModel - Full context optimization",
975
+ ],
976
+ style="success",
977
+ )
978
+
979
+ return 0
980
+
981
+ except Exception as e:
982
+ print_box(
983
+ "Error",
984
+ [
985
+ f"An error occurred: {type(e).__name__}",
986
+ "",
987
+ str(e)[:200],
988
+ "",
989
+ "Common issues:",
990
+ " - Invalid AWS credentials",
991
+ " - Bedrock not enabled in your AWS account",
992
+ " - Model not available in selected region",
993
+ " - Rate limiting from too many requests",
994
+ ],
995
+ style="error",
996
+ )
997
+ return 1
998
+
999
+
1000
+ if __name__ == "__main__":
1001
+ sys.exit(main())
headroom/integrations/strands/__init__.py ADDED
@@ -0,0 +1,88 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Strands Agents integration for Headroom SDK.
2
+
3
+ This module provides seamless integration with Strands Agents,
4
+ enabling automatic context optimization for Strands agents.
5
+
6
+ Components:
7
+ 1. HeadroomStrandsModel - Wraps any Strands model to apply Headroom transforms
8
+ 2. HeadroomHookProvider - Hook provider for Strands agents
9
+ 3. get_headroom_provider - Detects appropriate provider for a Strands model
10
+ 4. get_model_name_from_strands - Extracts model name from a Strands model
11
+
12
+ Example:
13
+ from strands import Agent
14
+ from strands.models import BedrockModel
15
+ from headroom.integrations.strands import HeadroomStrandsModel
16
+
17
+ # Wrap any Strands model
18
+ model = BedrockModel(model_id="anthropic.claude-3-5-sonnet-20241022-v2:0")
19
+ optimized_model = HeadroomStrandsModel(model)
20
+
21
+ # Use with agent
22
+ agent = Agent(model=optimized_model)
23
+ response = agent("Hello!")
24
+ """
25
+
26
+ from __future__ import annotations
27
+
28
+ import importlib.util
29
+ from typing import TYPE_CHECKING, Any
30
+
31
+ if TYPE_CHECKING:
32
+ from .hooks import HeadroomHookProvider
33
+ from .model import HeadroomStrandsModel, OptimizationMetrics, optimize_messages
34
+ from .providers import get_headroom_provider, get_model_name_from_strands
35
+
36
+
37
+ def strands_available() -> bool:
38
+ """Check if strands-agents is installed and available.
39
+
40
+ Returns:
41
+ True if strands-agents package is available, False otherwise.
42
+ """
43
+ return importlib.util.find_spec("strands") is not None
44
+
45
+
46
+ # Lazy imports to avoid import errors when strands is not installed
47
+ def __getattr__(name: str) -> Any:
48
+ """Lazy import of integration components."""
49
+ if name == "HeadroomHookProvider":
50
+ from .hooks import HeadroomHookProvider
51
+
52
+ return HeadroomHookProvider
53
+ elif name == "HeadroomStrandsModel":
54
+ from .model import HeadroomStrandsModel
55
+
56
+ return HeadroomStrandsModel
57
+ elif name == "OptimizationMetrics":
58
+ from .model import OptimizationMetrics
59
+
60
+ return OptimizationMetrics
61
+ elif name == "optimize_messages":
62
+ from .model import optimize_messages
63
+
64
+ return optimize_messages
65
+ elif name == "get_headroom_provider":
66
+ from .providers import get_headroom_provider
67
+
68
+ return get_headroom_provider
69
+ elif name == "get_model_name_from_strands":
70
+ from .providers import get_model_name_from_strands
71
+
72
+ return get_model_name_from_strands
73
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
74
+
75
+
76
+ __all__ = [
77
+ # Availability check
78
+ "strands_available",
79
+ # Hook provider
80
+ "HeadroomHookProvider",
81
+ # Model wrapper
82
+ "HeadroomStrandsModel",
83
+ "OptimizationMetrics",
84
+ "optimize_messages",
85
+ # Provider detection
86
+ "get_headroom_provider",
87
+ "get_model_name_from_strands",
88
+ ]
headroom/integrations/strands/hooks.py ADDED
@@ -0,0 +1,540 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Strands SDK hook provider for Headroom tool output compression.
2
+
3
+ This module provides HeadroomHookProvider, which implements Strands' HookProvider
4
+ interface to intercept tool outputs and compress them using Headroom's SmartCrusher.
5
+
6
+ Example:
7
+ from strands import Agent
8
+ from headroom.integrations.strands import HeadroomHookProvider
9
+
10
+ # Create the hook provider
11
+ hook_provider = HeadroomHookProvider(
12
+ compress_tool_outputs=True,
13
+ min_tokens_to_compress=100,
14
+ )
15
+
16
+ # Use with Strands agent
17
+ agent = Agent(hooks=[hook_provider])
18
+ response = agent("Search for documents about AI")
19
+
20
+ # Check compression metrics
21
+ print(f"Tokens saved: {hook_provider.total_tokens_saved}")
22
+ """
23
+
24
+ from __future__ import annotations
25
+
26
+ import json
27
+ import logging
28
+ import threading
29
+ from dataclasses import dataclass, field
30
+ from datetime import datetime, timezone
31
+ from typing import Any
32
+ from uuid import uuid4
33
+
34
+ # Strands imports - these are optional dependencies
35
+ try:
36
+ from strands.hooks import HookProvider, HookRegistry
37
+ from strands.hooks.events import AfterToolCallEvent, BeforeToolCallEvent
38
+ from strands.types.tools import ToolResult
39
+
40
+ STRANDS_AVAILABLE = True
41
+ except ImportError:
42
+ STRANDS_AVAILABLE = False
43
+ # Type stubs for when strands is not installed
44
+ HookProvider = object # type: ignore[misc,assignment]
45
+ HookRegistry = object # type: ignore[misc,assignment]
46
+ AfterToolCallEvent = object # type: ignore[misc,assignment]
47
+ BeforeToolCallEvent = object # type: ignore[misc,assignment]
48
+ ToolResult = dict # type: ignore[misc,assignment]
49
+
50
+ from headroom import HeadroomConfig
51
+ from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig
52
+
53
+ logger = logging.getLogger(__name__)
54
+
55
+
56
+ def _check_strands_available() -> None:
57
+ """Raise ImportError if Strands is not installed."""
58
+ if not STRANDS_AVAILABLE:
59
+ raise ImportError(
60
+ "Strands SDK is required for this integration. Install with: pip install strands-agents"
61
+ )
62
+
63
+
64
+ def strands_available() -> bool:
65
+ """Check if Strands SDK is installed.
66
+
67
+ Returns:
68
+ True if strands-agents package is available.
69
+ """
70
+ return STRANDS_AVAILABLE
71
+
72
+
73
+ @dataclass
74
+ class CompressionMetrics:
75
+ """Metrics from a single tool output compression."""
76
+
77
+ request_id: str
78
+ timestamp: datetime
79
+ tool_name: str
80
+ tool_use_id: str
81
+ tokens_before: int
82
+ tokens_after: int
83
+ tokens_saved: int
84
+ savings_percent: float
85
+ was_compressed: bool
86
+ skip_reason: str | None = None
87
+
88
+
89
+ @dataclass
90
+ class HeadroomHookProvider(HookProvider): # type: ignore[misc]
91
+ """Strands HookProvider that compresses tool outputs using Headroom.
92
+
93
+ This hook provider intercepts tool call results via AfterToolCallEvent
94
+ and applies Headroom's SmartCrusher to compress large outputs, reducing
95
+ token usage while preserving important information.
96
+
97
+ The compression is intelligent and preserves:
98
+ - Error items (containing error indicators)
99
+ - Anomalous values (statistical outliers)
100
+ - Items matching the user's query context
101
+ - First/last items for context
102
+ - Structural outliers (rare status values)
103
+
104
+ Attributes:
105
+ compress_tool_outputs: Whether to compress tool outputs.
106
+ min_tokens_to_compress: Minimum token count before compression is applied.
107
+ config: Headroom configuration.
108
+ preserve_errors: If True, never compress results with error status.
109
+ total_tokens_saved: Running total of tokens saved across all compressions.
110
+ metrics_history: List of CompressionMetrics from recent compressions.
111
+
112
+ Example:
113
+ from strands import Agent
114
+ from headroom.integrations.strands import HeadroomHookProvider
115
+
116
+ hook = HeadroomHookProvider(min_tokens_to_compress=50)
117
+ agent = Agent(hooks=[hook])
118
+
119
+ # After running agent tasks...
120
+ summary = hook.get_savings_summary()
121
+ print(f"Total saved: {summary['total_tokens_saved']} tokens")
122
+ """
123
+
124
+ compress_tool_outputs: bool = True
125
+ min_tokens_to_compress: int = 100
126
+ config: HeadroomConfig | None = field(default=None)
127
+ preserve_errors: bool = True
128
+
129
+ # Internal state (not part of dataclass comparison)
130
+ _crusher: SmartCrusher | None = field(default=None, repr=False, compare=False)
131
+ _metrics_history: list[CompressionMetrics] = field(
132
+ default_factory=list, repr=False, compare=False
133
+ )
134
+ _total_tokens_saved: int = field(default=0, repr=False, compare=False)
135
+ _lock: threading.Lock = field(default_factory=threading.Lock, repr=False, compare=False)
136
+ _initialized: bool = field(default=False, repr=False, compare=False)
137
+
138
+ def __post_init__(self) -> None:
139
+ """Initialize the hook provider after dataclass construction."""
140
+ _check_strands_available()
141
+
142
+ if self.config is None:
143
+ self.config = HeadroomConfig()
144
+
145
+ self._initialized = True
146
+ logger.debug(
147
+ "HeadroomHookProvider initialized: compress=%s, min_tokens=%d, preserve_errors=%s",
148
+ self.compress_tool_outputs,
149
+ self.min_tokens_to_compress,
150
+ self.preserve_errors,
151
+ )
152
+
153
+ @property
154
+ def crusher(self) -> SmartCrusher:
155
+ """Lazily initialize SmartCrusher (thread-safe).
156
+
157
+ Returns:
158
+ The SmartCrusher instance for compression.
159
+ """
160
+ if self._crusher is None:
161
+ with self._lock:
162
+ # Double-check after acquiring lock
163
+ if self._crusher is None:
164
+ # Use config from HeadroomConfig if available
165
+ if self.config and self.config.smart_crusher:
166
+ crusher_config = SmartCrusherConfig(
167
+ min_tokens_to_crush=self.min_tokens_to_compress,
168
+ max_items_after_crush=self.config.smart_crusher.max_items_after_crush,
169
+ )
170
+ else:
171
+ crusher_config = SmartCrusherConfig(
172
+ min_tokens_to_crush=self.min_tokens_to_compress
173
+ )
174
+ self._crusher = SmartCrusher(config=crusher_config)
175
+ logger.debug(
176
+ "SmartCrusher initialized with min_tokens=%d", self.min_tokens_to_compress
177
+ )
178
+ return self._crusher
179
+
180
+ @property
181
+ def total_tokens_saved(self) -> int:
182
+ """Total tokens saved across all compressions.
183
+
184
+ Returns:
185
+ Cumulative token savings.
186
+ """
187
+ return self._total_tokens_saved
188
+
189
+ @property
190
+ def metrics_history(self) -> list[CompressionMetrics]:
191
+ """History of compression metrics.
192
+
193
+ Returns:
194
+ Copy of the metrics history list.
195
+ """
196
+ return self._metrics_history.copy()
197
+
198
+ def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None:
199
+ """Register hooks with the Strands HookRegistry.
200
+
201
+ This method is called by Strands when the hook provider is added
202
+ to an Agent. It registers the compression handler for AfterToolCallEvent.
203
+
204
+ Args:
205
+ registry: The Strands HookRegistry to register hooks with.
206
+ """
207
+ if not self.compress_tool_outputs:
208
+ logger.debug("Tool output compression disabled, skipping hook registration")
209
+ return
210
+
211
+ # Register the after-tool-call hook for compression
212
+ registry.add_callback(AfterToolCallEvent, self._compress_tool_result)
213
+ logger.info(
214
+ "HeadroomHookProvider registered: compressing tool outputs >= %d tokens",
215
+ self.min_tokens_to_compress,
216
+ )
217
+
218
+ def _estimate_tokens(self, text: str) -> int:
219
+ """Estimate token count for text.
220
+
221
+ Uses a simple heuristic of ~4 characters per token, which is
222
+ reasonably accurate for English text and JSON content.
223
+
224
+ Args:
225
+ text: The text to estimate tokens for.
226
+
227
+ Returns:
228
+ Estimated token count.
229
+ """
230
+ if not text:
231
+ return 0
232
+ # ~4 characters per token is a reasonable estimate
233
+ return len(text) // 4
234
+
235
+ def _extract_text_content(self, result: ToolResult) -> str:
236
+ """Extract text content from a ToolResult.
237
+
238
+ Handles both text and JSON content types in the result.
239
+
240
+ Args:
241
+ result: The ToolResult to extract content from.
242
+
243
+ Returns:
244
+ String representation of the content.
245
+ """
246
+ content = result.get("content", [])
247
+ if not content:
248
+ return ""
249
+
250
+ text_parts = []
251
+ for item in content:
252
+ if isinstance(item, dict):
253
+ if "text" in item:
254
+ text_parts.append(str(item["text"]))
255
+ elif "json" in item:
256
+ try:
257
+ text_parts.append(json.dumps(item["json"], indent=None))
258
+ except (TypeError, ValueError):
259
+ text_parts.append(str(item["json"]))
260
+ elif isinstance(item, str):
261
+ text_parts.append(item)
262
+
263
+ return "\n".join(text_parts)
264
+
265
+ def _update_result_content(self, result: ToolResult, compressed_text: str) -> None:
266
+ """Update the result content with compressed text.
267
+
268
+ Modifies the result in place, preserving the original content structure
269
+ (text vs json) where possible.
270
+
271
+ Args:
272
+ result: The ToolResult to update (modified in place).
273
+ compressed_text: The compressed content to set.
274
+ """
275
+ content = result.get("content", [])
276
+
277
+ if not content:
278
+ # No existing content, create text content
279
+ result["content"] = [{"text": compressed_text}]
280
+ return
281
+
282
+ # Try to preserve original structure
283
+ first_item = content[0] if content else None
284
+
285
+ if isinstance(first_item, dict):
286
+ if "json" in first_item:
287
+ # Try to parse compressed text back to JSON
288
+ try:
289
+ parsed = json.loads(compressed_text)
290
+ result["content"] = [{"json": parsed}]
291
+ except (json.JSONDecodeError, ValueError):
292
+ # Fall back to text if not valid JSON
293
+ result["content"] = [{"text": compressed_text}]
294
+ else:
295
+ # Text content
296
+ result["content"] = [{"text": compressed_text}]
297
+ else:
298
+ # Unknown structure, use text
299
+ result["content"] = [{"text": compressed_text}]
300
+
301
+ def _compress_tool_result(self, event: AfterToolCallEvent) -> None:
302
+ """Compress tool result content if it exceeds the token threshold.
303
+
304
+ This is the main hook handler that intercepts AfterToolCallEvent
305
+ and applies SmartCrusher compression to large tool outputs.
306
+
307
+ Args:
308
+ event: The AfterToolCallEvent containing the tool result.
309
+ The result field is writable and modified in place.
310
+ """
311
+ request_id = str(uuid4())
312
+ result = event.result
313
+ tool_name = event.tool_use.get("name", "unknown")
314
+ tool_use_id = event.tool_use.get("toolUseId", "unknown")
315
+
316
+ # Check if compression should be skipped
317
+ skip_reason = self._should_skip_compression(result)
318
+ if skip_reason:
319
+ self._record_metrics(
320
+ request_id=request_id,
321
+ tool_name=tool_name,
322
+ tool_use_id=tool_use_id,
323
+ tokens_before=0,
324
+ tokens_after=0,
325
+ was_compressed=False,
326
+ skip_reason=skip_reason,
327
+ )
328
+ logger.debug(
329
+ "Skipping compression for tool %s (id=%s): %s",
330
+ tool_name,
331
+ tool_use_id,
332
+ skip_reason,
333
+ )
334
+ return
335
+
336
+ # Extract content and estimate tokens
337
+ original_text = self._extract_text_content(result)
338
+ tokens_before = self._estimate_tokens(original_text)
339
+
340
+ # Check minimum token threshold
341
+ if tokens_before < self.min_tokens_to_compress:
342
+ self._record_metrics(
343
+ request_id=request_id,
344
+ tool_name=tool_name,
345
+ tool_use_id=tool_use_id,
346
+ tokens_before=tokens_before,
347
+ tokens_after=tokens_before,
348
+ was_compressed=False,
349
+ skip_reason=f"below_threshold:{tokens_before}<{self.min_tokens_to_compress}",
350
+ )
351
+ logger.debug(
352
+ "Tool %s output below threshold (%d < %d tokens), skipping compression",
353
+ tool_name,
354
+ tokens_before,
355
+ self.min_tokens_to_compress,
356
+ )
357
+ return
358
+
359
+ # Apply compression
360
+ try:
361
+ crush_result = self.crusher.crush(content=original_text, query="")
362
+ compressed_text = crush_result.compressed
363
+ was_modified = crush_result.was_modified
364
+ except Exception as e:
365
+ # Compression failed, keep original
366
+ logger.warning(
367
+ "Compression failed for tool %s (id=%s): %s. Keeping original.",
368
+ tool_name,
369
+ tool_use_id,
370
+ str(e),
371
+ )
372
+ self._record_metrics(
373
+ request_id=request_id,
374
+ tool_name=tool_name,
375
+ tool_use_id=tool_use_id,
376
+ tokens_before=tokens_before,
377
+ tokens_after=tokens_before,
378
+ was_compressed=False,
379
+ skip_reason=f"compression_error:{type(e).__name__}",
380
+ )
381
+ return
382
+
383
+ tokens_after = self._estimate_tokens(compressed_text)
384
+
385
+ # Only update if compression actually reduced tokens
386
+ if was_modified and tokens_after < tokens_before:
387
+ self._update_result_content(result, compressed_text)
388
+ tokens_saved = tokens_before - tokens_after
389
+
390
+ self._record_metrics(
391
+ request_id=request_id,
392
+ tool_name=tool_name,
393
+ tool_use_id=tool_use_id,
394
+ tokens_before=tokens_before,
395
+ tokens_after=tokens_after,
396
+ was_compressed=True,
397
+ skip_reason=None,
398
+ )
399
+
400
+ logger.info(
401
+ "Compressed tool %s output: %d -> %d tokens (%.1f%% saved)",
402
+ tool_name,
403
+ tokens_before,
404
+ tokens_after,
405
+ (tokens_saved / tokens_before * 100) if tokens_before > 0 else 0,
406
+ )
407
+ else:
408
+ # Compression didn't help
409
+ self._record_metrics(
410
+ request_id=request_id,
411
+ tool_name=tool_name,
412
+ tool_use_id=tool_use_id,
413
+ tokens_before=tokens_before,
414
+ tokens_after=tokens_before,
415
+ was_compressed=False,
416
+ skip_reason="no_reduction",
417
+ )
418
+ logger.debug(
419
+ "Compression did not reduce tool %s output (%d tokens)",
420
+ tool_name,
421
+ tokens_before,
422
+ )
423
+
424
+ def _should_skip_compression(self, result: ToolResult) -> str | None:
425
+ """Check if compression should be skipped for this result.
426
+
427
+ Args:
428
+ result: The tool result to check.
429
+
430
+ Returns:
431
+ Skip reason string if should skip, None if should compress.
432
+ """
433
+ # Skip if compression is disabled
434
+ if not self.compress_tool_outputs:
435
+ return "compression_disabled"
436
+
437
+ # Skip error results if preserve_errors is True
438
+ if self.preserve_errors and result.get("status") == "error":
439
+ return "error_result_preserved"
440
+
441
+ # Skip empty results
442
+ content = result.get("content", [])
443
+ if not content:
444
+ return "empty_content"
445
+
446
+ return None
447
+
448
+ def _record_metrics(
449
+ self,
450
+ request_id: str,
451
+ tool_name: str,
452
+ tool_use_id: str,
453
+ tokens_before: int,
454
+ tokens_after: int,
455
+ was_compressed: bool,
456
+ skip_reason: str | None,
457
+ ) -> None:
458
+ """Record compression metrics (thread-safe).
459
+
460
+ Args:
461
+ request_id: Unique ID for this compression request.
462
+ tool_name: Name of the tool that was called.
463
+ tool_use_id: The toolUseId from the result.
464
+ tokens_before: Token count before compression.
465
+ tokens_after: Token count after compression.
466
+ was_compressed: Whether compression was actually applied.
467
+ skip_reason: Reason compression was skipped, if applicable.
468
+ """
469
+ tokens_saved = max(0, tokens_before - tokens_after)
470
+ savings_percent = (tokens_saved / tokens_before * 100) if tokens_before > 0 else 0.0
471
+
472
+ metrics = CompressionMetrics(
473
+ request_id=request_id,
474
+ timestamp=datetime.now(timezone.utc),
475
+ tool_name=tool_name,
476
+ tool_use_id=tool_use_id,
477
+ tokens_before=tokens_before,
478
+ tokens_after=tokens_after,
479
+ tokens_saved=tokens_saved,
480
+ savings_percent=savings_percent,
481
+ was_compressed=was_compressed,
482
+ skip_reason=skip_reason,
483
+ )
484
+
485
+ with self._lock:
486
+ self._metrics_history.append(metrics)
487
+ if was_compressed:
488
+ self._total_tokens_saved += tokens_saved
489
+
490
+ # Keep only last 100 metrics to bound memory
491
+ if len(self._metrics_history) > 100:
492
+ self._metrics_history = self._metrics_history[-100:]
493
+
494
+ def get_savings_summary(self) -> dict[str, Any]:
495
+ """Get summary of token savings across all compressions.
496
+
497
+ Returns:
498
+ Dictionary with compression statistics including:
499
+ - total_requests: Number of tool outputs processed
500
+ - compressed_requests: Number actually compressed
501
+ - total_tokens_saved: Cumulative tokens saved
502
+ - average_savings_percent: Mean compression ratio
503
+ - total_tokens_before: Sum of all input tokens
504
+ - total_tokens_after: Sum of all output tokens
505
+ """
506
+ if not self._metrics_history:
507
+ return {
508
+ "total_requests": 0,
509
+ "compressed_requests": 0,
510
+ "total_tokens_saved": 0,
511
+ "average_savings_percent": 0.0,
512
+ "total_tokens_before": 0,
513
+ "total_tokens_after": 0,
514
+ }
515
+
516
+ compressed_metrics = [m for m in self._metrics_history if m.was_compressed]
517
+
518
+ return {
519
+ "total_requests": len(self._metrics_history),
520
+ "compressed_requests": len(compressed_metrics),
521
+ "total_tokens_saved": self._total_tokens_saved,
522
+ "average_savings_percent": (
523
+ sum(m.savings_percent for m in compressed_metrics) / len(compressed_metrics)
524
+ if compressed_metrics
525
+ else 0.0
526
+ ),
527
+ "total_tokens_before": sum(m.tokens_before for m in self._metrics_history),
528
+ "total_tokens_after": sum(m.tokens_after for m in self._metrics_history),
529
+ }
530
+
531
+ def reset(self) -> None:
532
+ """Reset all tracked metrics (thread-safe).
533
+
534
+ Clears the metrics history and resets the total tokens saved counter.
535
+ Useful for starting fresh measurements or between test runs.
536
+ """
537
+ with self._lock:
538
+ self._metrics_history = []
539
+ self._total_tokens_saved = 0
540
+ logger.debug("HeadroomHookProvider metrics reset")
headroom/integrations/strands/model.py ADDED
@@ -0,0 +1,625 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Strands SDK model wrapper for Headroom optimization.
2
+
3
+ This module provides HeadroomStrandsModel, which wraps any Strands model
4
+ to apply Headroom context optimization before API calls.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import asyncio
10
+ import logging
11
+ import threading
12
+ from collections.abc import AsyncGenerator, AsyncIterable
13
+ from dataclasses import dataclass
14
+ from datetime import datetime, timezone
15
+ from typing import Any, TypeVar
16
+ from uuid import uuid4
17
+
18
+ # Strands imports - these are optional dependencies
19
+ try:
20
+ from strands.models import Model
21
+ from strands.types.content import Message, Messages, SystemContentBlock
22
+ from strands.types.streaming import StreamEvent
23
+ from strands.types.tools import ToolChoice, ToolSpec
24
+
25
+ STRANDS_AVAILABLE = True
26
+ except ImportError:
27
+ STRANDS_AVAILABLE = False
28
+ Model = object # type: ignore[misc,assignment]
29
+ Message = dict # type: ignore[misc,assignment]
30
+ Messages = list # type: ignore[misc,assignment]
31
+ StreamEvent = dict # type: ignore[misc,assignment]
32
+ ToolChoice = dict # type: ignore[misc,assignment]
33
+ ToolSpec = dict # type: ignore[misc,assignment]
34
+ SystemContentBlock = dict # type: ignore[misc,assignment]
35
+
36
+ T = TypeVar("T")
37
+
38
+ from headroom import HeadroomConfig # noqa: E402
39
+ from headroom.providers import OpenAIProvider # noqa: E402
40
+ from headroom.transforms import TransformPipeline # noqa: E402
41
+
42
+ from .providers import get_headroom_provider, get_model_name_from_strands # noqa: E402
43
+
44
+ logger = logging.getLogger(__name__)
45
+
46
+
47
+ def _check_strands_available() -> None:
48
+ """Raise ImportError if Strands SDK is not installed."""
49
+ if not STRANDS_AVAILABLE:
50
+ raise ImportError(
51
+ "Strands SDK is required for this integration. Install with: pip install strands-agents"
52
+ )
53
+
54
+
55
+ def strands_available() -> bool:
56
+ """Check if Strands SDK is installed."""
57
+ return STRANDS_AVAILABLE
58
+
59
+
60
+ @dataclass
61
+ class OptimizationMetrics:
62
+ """Metrics from a single optimization pass."""
63
+
64
+ request_id: str
65
+ timestamp: datetime
66
+ tokens_before: int
67
+ tokens_after: int
68
+ tokens_saved: int
69
+ savings_percent: float
70
+ transforms_applied: list[str]
71
+ model: str
72
+
73
+
74
+ class HeadroomStrandsModel(Model): # type: ignore[misc]
75
+ """Strands model wrapper that applies Headroom optimizations.
76
+
77
+ Wraps any Strands Model and automatically optimizes the context
78
+ before each API call. Works with any Strands-compatible model provider.
79
+
80
+ Example:
81
+ from strands import Agent
82
+ from strands.models.bedrock import BedrockModel
83
+ from headroom.integrations.strands import HeadroomStrandsModel
84
+
85
+ # Basic usage
86
+ model = BedrockModel(model_id="us.anthropic.claude-sonnet-4-20250514-v1:0")
87
+ optimized = HeadroomStrandsModel(wrapped_model=model)
88
+
89
+ # Use with agent
90
+ agent = Agent(model=optimized)
91
+ response = agent("Hello!")
92
+
93
+ # Access metrics
94
+ print(f"Saved {optimized.total_tokens_saved} tokens")
95
+
96
+ # With custom config
97
+ from headroom import HeadroomConfig
98
+ config = HeadroomConfig()
99
+ optimized = HeadroomStrandsModel(wrapped_model=model, config=config)
100
+
101
+ Attributes:
102
+ wrapped_model: The underlying Strands model
103
+ total_tokens_saved: Running total of tokens saved
104
+ metrics_history: List of OptimizationMetrics from recent calls
105
+ """
106
+
107
+ def __init__(
108
+ self,
109
+ wrapped_model: Any,
110
+ config: HeadroomConfig | None = None,
111
+ auto_detect_provider: bool = True,
112
+ ) -> None:
113
+ """Initialize HeadroomStrandsModel.
114
+
115
+ Args:
116
+ wrapped_model: The Strands model to wrap (e.g., BedrockModel, OpenAIModel)
117
+ config: Optional HeadroomConfig for optimization settings
118
+ auto_detect_provider: Whether to auto-detect the Headroom provider
119
+ based on the wrapped model type. Default True.
120
+ """
121
+ _check_strands_available()
122
+
123
+ if wrapped_model is None:
124
+ raise ValueError("wrapped_model cannot be None")
125
+
126
+ self.wrapped_model = wrapped_model
127
+ self.headroom_config = config or HeadroomConfig()
128
+ self.auto_detect_provider = auto_detect_provider
129
+
130
+ # Internal state
131
+ self._metrics_history: list[OptimizationMetrics] = []
132
+ self._total_tokens_saved: int = 0
133
+ self._pipeline: TransformPipeline | None = None
134
+ self._headroom_provider: Any = None
135
+ self._lock = threading.Lock()
136
+
137
+ @property
138
+ def config(self) -> Any:
139
+ """Forward config access to wrapped model (required by Strands Agent)."""
140
+ return self.wrapped_model.config
141
+
142
+ @property
143
+ def pipeline(self) -> TransformPipeline:
144
+ """Lazily initialize TransformPipeline (thread-safe)."""
145
+ if self._pipeline is None:
146
+ with self._lock:
147
+ # Double-check after acquiring lock
148
+ if self._pipeline is None:
149
+ if self.auto_detect_provider:
150
+ self._headroom_provider = get_headroom_provider(self.wrapped_model)
151
+ logger.debug(
152
+ f"Auto-detected provider: {self._headroom_provider.__class__.__name__}"
153
+ )
154
+ else:
155
+ self._headroom_provider = OpenAIProvider()
156
+ self._pipeline = TransformPipeline(
157
+ config=self.headroom_config,
158
+ provider=self._headroom_provider,
159
+ )
160
+ return self._pipeline
161
+
162
+ @property
163
+ def total_tokens_saved(self) -> int:
164
+ """Total tokens saved across all calls."""
165
+ return self._total_tokens_saved
166
+
167
+ @property
168
+ def metrics_history(self) -> list[OptimizationMetrics]:
169
+ """History of optimization metrics."""
170
+ return self._metrics_history.copy()
171
+
172
+ def _convert_messages_to_openai(self, messages: list[Any]) -> list[dict[str, Any]]:
173
+ """Convert Strands messages to OpenAI format for Headroom.
174
+
175
+ Strands uses dict-based messages similar to OpenAI format:
176
+ - {"role": "user", "content": "..."}
177
+ - {"role": "assistant", "content": "...", "tool_calls": [...]}
178
+ - {"role": "tool", "content": "...", "tool_call_id": "..."}
179
+
180
+ Args:
181
+ messages: List of Strands messages (typically dicts or Message objects)
182
+
183
+ Returns:
184
+ List of messages in OpenAI dict format
185
+ """
186
+ result = []
187
+ for msg in messages:
188
+ # Handle dict format (most common in Strands)
189
+ if isinstance(msg, dict):
190
+ entry: dict[str, Any] = {
191
+ "role": msg.get("role", "user"),
192
+ }
193
+
194
+ # Handle content
195
+ content = msg.get("content")
196
+ if content is None:
197
+ entry["content"] = ""
198
+ elif isinstance(content, list):
199
+ # Content blocks - preserve structure
200
+ entry["content"] = content
201
+ else:
202
+ entry["content"] = content
203
+
204
+ # Handle tool calls
205
+ if "tool_calls" in msg and msg["tool_calls"]:
206
+ entry["tool_calls"] = msg["tool_calls"]
207
+
208
+ # Handle tool call ID for tool responses
209
+ if "tool_call_id" in msg and msg["tool_call_id"]:
210
+ entry["tool_call_id"] = msg["tool_call_id"]
211
+
212
+ # Handle name field (for tool messages)
213
+ if "name" in msg and msg["name"]:
214
+ entry["name"] = msg["name"]
215
+
216
+ result.append(entry)
217
+
218
+ # Handle Strands Message objects (if they have role/content attrs)
219
+ elif hasattr(msg, "role") and hasattr(msg, "content"):
220
+ entry = {
221
+ "role": msg.role,
222
+ }
223
+
224
+ content = msg.content
225
+ if content is None:
226
+ entry["content"] = ""
227
+ elif isinstance(content, list):
228
+ entry["content"] = content
229
+ else:
230
+ entry["content"] = content
231
+
232
+ if hasattr(msg, "tool_calls") and msg.tool_calls:
233
+ entry["tool_calls"] = msg.tool_calls
234
+ if hasattr(msg, "tool_call_id") and msg.tool_call_id:
235
+ entry["tool_call_id"] = msg.tool_call_id
236
+ if hasattr(msg, "name") and msg.name:
237
+ entry["name"] = msg.name
238
+
239
+ result.append(entry)
240
+
241
+ else:
242
+ # Fallback: convert to string
243
+ content = str(msg) if msg is not None else ""
244
+ result.append({"role": "user", "content": content})
245
+
246
+ return result
247
+
248
+ def _convert_messages_from_openai(
249
+ self, messages: list[dict[str, Any]], original_messages: list[Any]
250
+ ) -> list[dict[str, Any]]:
251
+ """Convert OpenAI format messages back to Strands format.
252
+
253
+ Since Strands uses dict-based messages similar to OpenAI,
254
+ this is largely a passthrough, but ensures proper structure.
255
+
256
+ Args:
257
+ messages: The optimized messages in OpenAI dict format
258
+ original_messages: The original Strands messages (for reference)
259
+
260
+ Returns:
261
+ List of messages in Strands dict format
262
+ """
263
+ result = []
264
+ for msg in messages:
265
+ entry: dict[str, Any] = {
266
+ "role": msg.get("role", "user"),
267
+ }
268
+
269
+ # Handle content
270
+ content = msg.get("content")
271
+ if content is not None:
272
+ entry["content"] = content
273
+
274
+ # Preserve tool-related fields
275
+ if "tool_calls" in msg and msg["tool_calls"]:
276
+ entry["tool_calls"] = msg["tool_calls"]
277
+ if "tool_call_id" in msg and msg["tool_call_id"]:
278
+ entry["tool_call_id"] = msg["tool_call_id"]
279
+ if "name" in msg and msg["name"]:
280
+ entry["name"] = msg["name"]
281
+
282
+ result.append(entry)
283
+
284
+ return result
285
+
286
+ def _optimize_messages(
287
+ self, messages: list[Any]
288
+ ) -> tuple[list[dict[str, Any]], OptimizationMetrics]:
289
+ """Apply Headroom optimization to messages.
290
+
291
+ Thread-safe with fallback on pipeline errors.
292
+
293
+ Args:
294
+ messages: List of Strands messages to optimize
295
+
296
+ Returns:
297
+ Tuple of (optimized_messages, metrics)
298
+ """
299
+ request_id = str(uuid4())
300
+
301
+ # Convert to OpenAI format
302
+ openai_messages = self._convert_messages_to_openai(messages)
303
+
304
+ # Handle empty messages gracefully
305
+ if not openai_messages:
306
+ metrics = OptimizationMetrics(
307
+ request_id=request_id,
308
+ timestamp=datetime.now(timezone.utc),
309
+ tokens_before=0,
310
+ tokens_after=0,
311
+ tokens_saved=0,
312
+ savings_percent=0,
313
+ transforms_applied=[],
314
+ model=get_model_name_from_strands(self.wrapped_model),
315
+ )
316
+ return [], metrics
317
+
318
+ # Get model name from wrapped model
319
+ model = get_model_name_from_strands(self.wrapped_model)
320
+
321
+ # Ensure pipeline is initialized
322
+ _ = self.pipeline
323
+
324
+ # Get model context limit
325
+ model_limit = (
326
+ self._headroom_provider.get_context_limit(model) if self._headroom_provider else 128000
327
+ )
328
+
329
+ try:
330
+ # Apply Headroom transforms via pipeline
331
+ result = self.pipeline.apply(
332
+ messages=openai_messages,
333
+ model=model,
334
+ model_limit=model_limit,
335
+ )
336
+ optimized = result.messages
337
+ tokens_before = result.tokens_before
338
+ tokens_after = result.tokens_after
339
+ transforms_applied = result.transforms_applied
340
+ except (
341
+ ValueError,
342
+ TypeError,
343
+ AttributeError,
344
+ RuntimeError,
345
+ KeyError,
346
+ IndexError,
347
+ ImportError,
348
+ OSError,
349
+ ) as e:
350
+ # Fallback to original messages on pipeline error
351
+ logger.warning(
352
+ f"Headroom optimization failed, using original messages: {type(e).__name__}: {e}"
353
+ )
354
+ optimized = openai_messages
355
+ # Estimate token count (rough approximation: ~4 chars/token)
356
+ tokens_before = sum(len(str(m.get("content", ""))) // 4 for m in openai_messages)
357
+ tokens_after = tokens_before
358
+ transforms_applied = ["fallback:error"]
359
+
360
+ # Create metrics
361
+ tokens_saved = max(0, tokens_before - tokens_after)
362
+ metrics = OptimizationMetrics(
363
+ request_id=request_id,
364
+ timestamp=datetime.now(timezone.utc),
365
+ tokens_before=tokens_before,
366
+ tokens_after=tokens_after,
367
+ tokens_saved=tokens_saved,
368
+ savings_percent=(tokens_saved / tokens_before * 100 if tokens_before > 0 else 0),
369
+ transforms_applied=transforms_applied,
370
+ model=model,
371
+ )
372
+
373
+ # Track metrics (thread-safe)
374
+ with self._lock:
375
+ self._metrics_history.append(metrics)
376
+ self._total_tokens_saved += metrics.tokens_saved
377
+
378
+ # Keep only last 100 metrics
379
+ if len(self._metrics_history) > 100:
380
+ self._metrics_history = self._metrics_history[-100:]
381
+
382
+ # Convert back to Strands format
383
+ optimized_messages = self._convert_messages_from_openai(optimized, messages)
384
+
385
+ return optimized_messages, metrics
386
+
387
+ async def stream(
388
+ self,
389
+ messages: Messages,
390
+ tool_specs: list[ToolSpec] | None = None,
391
+ system_prompt: str | None = None,
392
+ *,
393
+ tool_choice: ToolChoice | None = None,
394
+ system_prompt_content: list[SystemContentBlock] | None = None,
395
+ invocation_state: dict[str, Any] | None = None,
396
+ **kwargs: Any,
397
+ ) -> AsyncIterable[StreamEvent]:
398
+ """Stream response with Headroom optimization.
399
+
400
+ This is the main method required by Strands Model interface.
401
+ Optimizes messages before delegating to the wrapped model's stream method.
402
+
403
+ Args:
404
+ messages: List of messages to send to the model
405
+ tool_specs: Optional list of tool specifications
406
+ system_prompt: Optional system prompt string
407
+ tool_choice: Optional tool choice configuration
408
+ system_prompt_content: Optional list of system content blocks
409
+ invocation_state: Optional invocation state dictionary
410
+ **kwargs: Additional arguments passed to the wrapped model
411
+
412
+ Yields:
413
+ Streaming events from the wrapped model
414
+ """
415
+ # Run optimization in executor (CPU-bound)
416
+ loop = asyncio.get_running_loop()
417
+ optimized_messages, metrics = await loop.run_in_executor(
418
+ None, self._optimize_messages, messages
419
+ )
420
+
421
+ logger.info(
422
+ f"Headroom optimized (stream): {metrics.tokens_before} -> "
423
+ f"{metrics.tokens_after} tokens ({metrics.savings_percent:.1f}% saved)"
424
+ )
425
+
426
+ # Delegate to wrapped model's stream method with all parameters
427
+ async for event in self.wrapped_model.stream(
428
+ optimized_messages,
429
+ tool_specs=tool_specs,
430
+ system_prompt=system_prompt,
431
+ tool_choice=tool_choice,
432
+ system_prompt_content=system_prompt_content,
433
+ invocation_state=invocation_state,
434
+ **kwargs,
435
+ ):
436
+ yield event
437
+
438
+ def get_config(self) -> Any:
439
+ """Get the configuration of the wrapped model.
440
+
441
+ Returns:
442
+ The model configuration from the wrapped model.
443
+ """
444
+ return self.wrapped_model.get_config()
445
+
446
+ def update_config(self, **model_config: Any) -> None:
447
+ """Update the configuration of the wrapped model.
448
+
449
+ Args:
450
+ **model_config: Configuration options to update on the wrapped model.
451
+ """
452
+ self.wrapped_model.update_config(**model_config)
453
+
454
+ async def structured_output(
455
+ self,
456
+ output_model: type[T],
457
+ prompt: Messages,
458
+ system_prompt: str | None = None,
459
+ **kwargs: Any,
460
+ ) -> AsyncGenerator[dict[str, T | Any], None]:
461
+ """Generate structured output with Headroom optimization.
462
+
463
+ Optimizes the prompt messages before delegating to the wrapped model's
464
+ structured_output method.
465
+
466
+ Args:
467
+ output_model: The type/schema for the structured output
468
+ prompt: List of prompt messages
469
+ system_prompt: Optional system prompt
470
+ **kwargs: Additional arguments passed to the wrapped model
471
+
472
+ Yields:
473
+ Structured output events from the wrapped model
474
+ """
475
+ # Run optimization in executor (CPU-bound)
476
+ loop = asyncio.get_running_loop()
477
+ optimized_prompt, metrics = await loop.run_in_executor(
478
+ None, self._optimize_messages, prompt
479
+ )
480
+
481
+ logger.info(
482
+ f"Headroom optimized (structured_output): {metrics.tokens_before} -> "
483
+ f"{metrics.tokens_after} tokens ({metrics.savings_percent:.1f}% saved)"
484
+ )
485
+
486
+ # Delegate to wrapped model
487
+ async for event in self.wrapped_model.structured_output(
488
+ output_model, optimized_prompt, system_prompt=system_prompt, **kwargs
489
+ ):
490
+ yield event
491
+
492
+ def get_savings_summary(self) -> dict[str, Any]:
493
+ """Get summary of token savings."""
494
+ if not self._metrics_history:
495
+ return {
496
+ "total_requests": 0,
497
+ "total_tokens_saved": 0,
498
+ "average_savings_percent": 0,
499
+ "total_tokens_before": 0,
500
+ "total_tokens_after": 0,
501
+ }
502
+
503
+ return {
504
+ "total_requests": len(self._metrics_history),
505
+ "total_tokens_saved": self._total_tokens_saved,
506
+ "average_savings_percent": sum(m.savings_percent for m in self._metrics_history)
507
+ / len(self._metrics_history),
508
+ "total_tokens_before": sum(m.tokens_before for m in self._metrics_history),
509
+ "total_tokens_after": sum(m.tokens_after for m in self._metrics_history),
510
+ }
511
+
512
+ def reset(self) -> None:
513
+ """Reset all tracked metrics (thread-safe).
514
+
515
+ Clears the metrics history and resets the total tokens saved counter.
516
+ Useful for starting fresh measurements or between test runs.
517
+ """
518
+ with self._lock:
519
+ self._metrics_history = []
520
+ self._total_tokens_saved = 0
521
+
522
+ # =========================================================================
523
+ # Forward attribute access to wrapped model for compatibility
524
+ # =========================================================================
525
+
526
+ def __getattr__(self, name: str) -> Any:
527
+ """Forward attribute access to wrapped model."""
528
+ # Avoid infinite recursion for our own attributes
529
+ if name in (
530
+ "wrapped_model",
531
+ "config",
532
+ "auto_detect_provider",
533
+ "_metrics_history",
534
+ "_total_tokens_saved",
535
+ "_pipeline",
536
+ "_headroom_provider",
537
+ "_lock",
538
+ "pipeline",
539
+ "total_tokens_saved",
540
+ "metrics_history",
541
+ ):
542
+ raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'")
543
+ return getattr(self.wrapped_model, name)
544
+
545
+
546
+ def optimize_messages(
547
+ messages: list[Any],
548
+ config: HeadroomConfig | None = None,
549
+ model: str = "gpt-4o",
550
+ ) -> tuple[list[dict[str, Any]], dict[str, Any]]:
551
+ """Standalone function to optimize Strands messages.
552
+
553
+ Use this for manual optimization when you need fine-grained control.
554
+
555
+ Args:
556
+ messages: List of Strands messages (dicts)
557
+ config: HeadroomConfig for optimization settings
558
+ model: Model name for token estimation
559
+
560
+ Returns:
561
+ Tuple of (optimized_messages, metrics_dict)
562
+
563
+ Example:
564
+ from headroom.integrations.strands import optimize_messages
565
+
566
+ messages = [
567
+ {"role": "system", "content": "You are helpful."},
568
+ {"role": "user", "content": "What is 2+2?"},
569
+ ]
570
+
571
+ optimized, metrics = optimize_messages(messages)
572
+ print(f"Saved {metrics['tokens_saved']} tokens")
573
+ """
574
+ _check_strands_available()
575
+
576
+ config = config or HeadroomConfig()
577
+ provider = OpenAIProvider()
578
+ pipeline = TransformPipeline(config=config, provider=provider)
579
+
580
+ # Convert to OpenAI format (Strands uses similar format)
581
+ openai_messages = []
582
+ for msg in messages:
583
+ if isinstance(msg, dict):
584
+ entry: dict[str, Any] = {
585
+ "role": msg.get("role", "user"),
586
+ "content": msg.get("content", ""),
587
+ }
588
+ if "tool_calls" in msg and msg["tool_calls"]:
589
+ entry["tool_calls"] = msg["tool_calls"]
590
+ if "tool_call_id" in msg and msg["tool_call_id"]:
591
+ entry["tool_call_id"] = msg["tool_call_id"]
592
+ openai_messages.append(entry)
593
+ elif hasattr(msg, "role") and hasattr(msg, "content"):
594
+ entry = {"role": msg.role, "content": msg.content or ""}
595
+ if hasattr(msg, "tool_calls") and msg.tool_calls:
596
+ entry["tool_calls"] = msg.tool_calls
597
+ if hasattr(msg, "tool_call_id") and msg.tool_call_id:
598
+ entry["tool_call_id"] = msg.tool_call_id
599
+ openai_messages.append(entry)
600
+ else:
601
+ openai_messages.append({"role": "user", "content": str(msg)})
602
+
603
+ # Get model context limit
604
+ model_limit = provider.get_context_limit(model)
605
+
606
+ # Apply transforms
607
+ result = pipeline.apply(
608
+ messages=openai_messages,
609
+ model=model,
610
+ model_limit=model_limit,
611
+ )
612
+
613
+ metrics = {
614
+ "tokens_before": result.tokens_before,
615
+ "tokens_after": result.tokens_after,
616
+ "tokens_saved": result.tokens_before - result.tokens_after,
617
+ "savings_percent": (
618
+ (result.tokens_before - result.tokens_after) / result.tokens_before * 100
619
+ if result.tokens_before > 0
620
+ else 0
621
+ ),
622
+ "transforms_applied": result.transforms_applied,
623
+ }
624
+
625
+ return result.messages, metrics
headroom/integrations/strands/providers.py ADDED
@@ -0,0 +1,166 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Provider detection for Strands models.
2
+
3
+ Automatically detects the correct Headroom provider based on the Strands model type.
4
+ """
5
+
6
+ from __future__ import annotations
7
+
8
+ import logging
9
+ from typing import Any
10
+
11
+ from headroom.providers import (
12
+ AnthropicProvider,
13
+ GoogleProvider,
14
+ OpenAIProvider,
15
+ )
16
+ from headroom.providers.base import Provider
17
+
18
+ logger = logging.getLogger(__name__)
19
+
20
+ # Mapping from Strands model class names to Headroom providers
21
+ _STRANDS_MODEL_PROVIDERS: dict[str, type[Provider]] = {
22
+ # Bedrock models (primarily Claude via Bedrock)
23
+ "BedrockModel": AnthropicProvider,
24
+ # Anthropic models (direct API)
25
+ "AnthropicModel": AnthropicProvider,
26
+ # OpenAI models
27
+ "OpenAIModel": OpenAIProvider,
28
+ # LiteLLM (uses OpenAI-compatible interface)
29
+ "LiteLLMModel": OpenAIProvider,
30
+ # Ollama (uses OpenAI-compatible interface)
31
+ "OllamaModel": OpenAIProvider,
32
+ # Google Gemini models
33
+ "GeminiModel": GoogleProvider,
34
+ # Writer models (uses OpenAI-compatible interface)
35
+ "WriterModel": OpenAIProvider,
36
+ }
37
+
38
+
39
+ def get_headroom_provider(model: Any) -> Provider:
40
+ """Get the appropriate Headroom provider for a Strands model.
41
+
42
+ Detection strategy:
43
+ 1. Check model class name against known Strands model types
44
+ 2. Check for provider hints in model attributes
45
+ 3. Fall back to OpenAI provider (most compatible)
46
+
47
+ Args:
48
+ model: A Strands model instance (BedrockModel, AnthropicModel, etc.)
49
+
50
+ Returns:
51
+ Appropriate Headroom Provider instance.
52
+
53
+ Example:
54
+ from strands.models import BedrockModel
55
+ from headroom.integrations.strands.providers import get_headroom_provider
56
+
57
+ model = BedrockModel(model_id="anthropic.claude-3-5-sonnet-20241022-v2:0")
58
+ provider = get_headroom_provider(model) # Returns AnthropicProvider
59
+ """
60
+ # Strategy 1: Class name matching
61
+ class_name = model.__class__.__name__
62
+ if class_name in _STRANDS_MODEL_PROVIDERS:
63
+ provider_class = _STRANDS_MODEL_PROVIDERS[class_name]
64
+ logger.debug(f"Detected provider {provider_class.__name__} from class {class_name}")
65
+ return provider_class()
66
+
67
+ # Strategy 2: Check module path
68
+ module_path = model.__class__.__module__
69
+ if "anthropic" in module_path.lower():
70
+ logger.debug(f"Detected AnthropicProvider from module {module_path}")
71
+ return AnthropicProvider()
72
+ elif "bedrock" in module_path.lower():
73
+ logger.debug(f"Detected AnthropicProvider from module {module_path}")
74
+ return AnthropicProvider()
75
+ elif "google" in module_path.lower() or "gemini" in module_path.lower():
76
+ logger.debug(f"Detected GoogleProvider from module {module_path}")
77
+ return GoogleProvider()
78
+ elif "openai" in module_path.lower() or "litellm" in module_path.lower():
79
+ logger.debug(f"Detected OpenAIProvider from module {module_path}")
80
+ return OpenAIProvider()
81
+
82
+ # Strategy 3: Check model ID/name for hints
83
+ model_id = _extract_model_id(model)
84
+ if model_id:
85
+ model_id_lower = model_id.lower()
86
+ if "claude" in model_id_lower or "anthropic" in model_id_lower:
87
+ logger.debug(f"Detected AnthropicProvider from model ID {model_id}")
88
+ return AnthropicProvider()
89
+ elif "gemini" in model_id_lower:
90
+ logger.debug(f"Detected GoogleProvider from model ID {model_id}")
91
+ return GoogleProvider()
92
+ elif "gpt" in model_id_lower or "o1" in model_id_lower or "o3" in model_id_lower:
93
+ logger.debug(f"Detected OpenAIProvider from model ID {model_id}")
94
+ return OpenAIProvider()
95
+
96
+ # Strategy 4: Default fallback
97
+ logger.warning(
98
+ f"Unknown Strands model class '{class_name}', defaulting to OpenAIProvider. "
99
+ "Token counting may be inaccurate."
100
+ )
101
+ return OpenAIProvider()
102
+
103
+
104
+ def _extract_model_id(model: Any) -> str:
105
+ """Extract model ID from a Strands model using various attribute names.
106
+
107
+ Args:
108
+ model: A Strands model instance
109
+
110
+ Returns:
111
+ Model ID string or empty string if not found
112
+ """
113
+ # Try common attribute names used by Strands models
114
+ for attr in ["model_id", "model", "model_name", "id"]:
115
+ value = getattr(model, attr, None)
116
+ if value and isinstance(value, str):
117
+ return str(value)
118
+
119
+ # Try to get from config if available (config can be dict or object)
120
+ config = getattr(model, "config", None)
121
+ if config:
122
+ for attr in ["model_id", "model", "model_name"]:
123
+ # Handle dict-style config (Strands uses this)
124
+ if isinstance(config, dict):
125
+ value = config.get(attr)
126
+ else:
127
+ value = getattr(config, attr, None)
128
+ if value and isinstance(value, str):
129
+ return str(value)
130
+
131
+ # Try get_config() method (Strands Model interface)
132
+ if hasattr(model, "get_config"):
133
+ try:
134
+ config_dict = model.get_config()
135
+ if isinstance(config_dict, dict):
136
+ for attr in ["model_id", "model", "model_name"]:
137
+ value = config_dict.get(attr)
138
+ if value and isinstance(value, str):
139
+ return str(value)
140
+ except Exception:
141
+ pass
142
+
143
+ return ""
144
+
145
+
146
+ def get_model_name_from_strands(model: Any) -> str:
147
+ """Extract the model name/ID from a Strands model.
148
+
149
+ Args:
150
+ model: A Strands model instance
151
+
152
+ Returns:
153
+ Model name string (e.g., "anthropic.claude-3-5-sonnet-20241022-v2:0")
154
+ """
155
+ model_id = _extract_model_id(model)
156
+ if model_id:
157
+ return str(model_id)
158
+
159
+ # Fallback with warning
160
+ class_name = model.__class__.__name__
161
+ logger.warning(
162
+ f"Could not extract model name from {class_name} (no 'model_id', 'model', "
163
+ f"'model_name', or 'id' attribute). Defaulting to 'gpt-4o'. "
164
+ "Token counting may be inaccurate."
165
+ )
166
+ return "gpt-4o"
pyproject.toml CHANGED
@@ -90,6 +90,10 @@ code = [
90
  agno = [
91
  "agno>=1.0.0",
92
  ]
 
 
 
 
93
  # Voice filler detection (training and inference)
94
  voice = [
95
  "onnxruntime>=1.16.0", # Fast CPU inference
 
90
  agno = [
91
  "agno>=1.0.0",
92
  ]
93
+ # AWS Strands Agents SDK integration
94
+ strands = [
95
+ "strands-agents>=0.1.0",
96
+ ]
97
  # Voice filler detection (training and inference)
98
  voice = [
99
  "onnxruntime>=1.16.0", # Fast CPU inference
tests/integrations/test_strands/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ """Tests for Strands Agents SDK integration with Headroom."""
tests/integrations/test_strands/test_hooks.py ADDED
@@ -0,0 +1,545 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Real-world integration tests for Strands HeadroomHookProvider.
2
+
3
+ These tests use actual AWS Bedrock API calls with real credentials.
4
+ NO MOCKS - all tests hit the real Bedrock API.
5
+
6
+ Skip in CI if AWS credentials are not available.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import json
12
+ import os
13
+
14
+ import pytest
15
+
16
+ # Check for AWS credentials availability
17
+ SKIP_BEDROCK = not (
18
+ os.environ.get("AWS_ACCESS_KEY_ID")
19
+ or os.environ.get("AWS_PROFILE")
20
+ or os.path.exists(os.path.expanduser("~/.aws/credentials"))
21
+ )
22
+
23
+ # Check if strands-agents is installed
24
+ try:
25
+ from strands import Agent, tool
26
+ from strands.models import BedrockModel
27
+
28
+ STRANDS_AVAILABLE = True
29
+ except ImportError:
30
+ STRANDS_AVAILABLE = False
31
+
32
+ # Provide a no-op decorator when strands is not installed
33
+ def tool(fn):
34
+ return fn
35
+
36
+ Agent = None # type: ignore
37
+ BedrockModel = None # type: ignore
38
+
39
+ # Skip all tests if dependencies not available
40
+ pytestmark = [
41
+ pytest.mark.skipif(SKIP_BEDROCK, reason="AWS credentials not available"),
42
+ pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed"),
43
+ ]
44
+
45
+
46
+ # ============================================================================
47
+ # Test Tools - Generate realistic verbose data for compression testing
48
+ # These are defined with @tool decorator for use when strands is installed.
49
+ # When strands is not installed, the no-op decorator ensures import succeeds.
50
+ # ============================================================================
51
+
52
+
53
+ @tool
54
+ def search_logs(query: str, limit: int = 100) -> str:
55
+ """Search application logs. Returns JSON array of log entries.
56
+
57
+ Args:
58
+ query: Search query to find in logs
59
+ limit: Maximum number of log entries to return
60
+
61
+ Returns:
62
+ JSON array of log entry objects
63
+ """
64
+ # Generate realistic verbose log data that should be compressed
65
+ logs = [
66
+ {
67
+ "timestamp": f"2024-01-{(i % 28) + 1:02d}T{10 + (i % 12):02d}:00:00Z",
68
+ "level": ["INFO", "DEBUG", "WARN", "ERROR"][i % 4],
69
+ "service": ["api-gateway", "auth-service", "data-processor", "cache-service"][i % 4],
70
+ "message": f"Request processed successfully - latency={50 + i}ms, query={query}",
71
+ "request_id": f"req-{i:06d}-{hash(query) % 10000:04d}",
72
+ "status_code": [200, 201, 400, 500][i % 4],
73
+ "user_agent": "Mozilla/5.0 (compatible; TestBot/1.0)",
74
+ "ip_address": f"192.168.{i % 256}.{(i * 7) % 256}",
75
+ "trace_id": f"trace-{i:08x}",
76
+ "span_id": f"span-{i:04x}",
77
+ "duration_ms": 50 + (i * 3) % 200,
78
+ "memory_mb": 128 + (i * 5) % 512,
79
+ "cpu_percent": 10 + (i * 2) % 80,
80
+ }
81
+ for i in range(limit)
82
+ ]
83
+ return json.dumps(logs, indent=2)
84
+
85
+
86
+ @tool
87
+ def get_small_status() -> str:
88
+ """Get a small status response that should NOT be compressed.
89
+
90
+ Returns:
91
+ Small JSON status object
92
+ """
93
+ return json.dumps({"status": "healthy", "uptime_seconds": 12345, "version": "1.2.3"})
94
+
95
+
96
+ @tool
97
+ def get_error_data() -> str:
98
+ """Get error information. Error results should NOT be compressed.
99
+
100
+ Returns:
101
+ Error information (but not as a tool error)
102
+ """
103
+ return json.dumps(
104
+ {
105
+ "errors": [
106
+ {"code": "E001", "message": "Connection timeout"},
107
+ {"code": "E002", "message": "Authentication failed"},
108
+ ],
109
+ "timestamp": "2024-01-15T10:00:00Z",
110
+ }
111
+ )
112
+
113
+
114
+ @tool
115
+ def fetch_user_data(user_id: str) -> str:
116
+ """Fetch detailed user data. Returns large JSON payload.
117
+
118
+ Args:
119
+ user_id: The user ID to fetch data for
120
+
121
+ Returns:
122
+ Large JSON object with user details
123
+ """
124
+ # Generate a large user profile that should trigger compression
125
+ activities = [
126
+ {
127
+ "activity_id": f"act-{i:06d}",
128
+ "type": ["login", "purchase", "view", "share"][i % 4],
129
+ "timestamp": f"2024-01-{(i % 28) + 1:02d}T{10 + (i % 12):02d}:30:00Z",
130
+ "details": {
131
+ "ip": f"10.0.{i % 256}.{(i * 3) % 256}",
132
+ "device": ["desktop", "mobile", "tablet"][i % 3],
133
+ "browser": ["Chrome", "Firefox", "Safari"][i % 3],
134
+ "duration_seconds": 30 + i * 5,
135
+ "page_views": 1 + i % 10,
136
+ },
137
+ "metadata": {
138
+ "session_id": f"sess-{i:08x}",
139
+ "referrer": f"https://example.com/page/{i}",
140
+ "utm_source": ["google", "facebook", "twitter", "email"][i % 4],
141
+ },
142
+ }
143
+ for i in range(50)
144
+ ]
145
+
146
+ return json.dumps(
147
+ {
148
+ "user_id": user_id,
149
+ "profile": {
150
+ "name": "Test User",
151
+ "email": f"{user_id}@example.com",
152
+ "created_at": "2023-01-01T00:00:00Z",
153
+ },
154
+ "activities": activities,
155
+ },
156
+ indent=2,
157
+ )
158
+
159
+
160
+ @tool
161
+ def simple_calculator(a: int, b: int, operation: str) -> str:
162
+ """Simple calculator for basic operations.
163
+
164
+ Args:
165
+ a: First number
166
+ b: Second number
167
+ operation: One of 'add', 'subtract', 'multiply', 'divide'
168
+
169
+ Returns:
170
+ The result of the operation
171
+ """
172
+ if operation == "add":
173
+ result = a + b
174
+ elif operation == "subtract":
175
+ result = a - b
176
+ elif operation == "multiply":
177
+ result = a * b
178
+ elif operation == "divide":
179
+ result = a / b if b != 0 else "undefined"
180
+ else:
181
+ result = "unknown operation"
182
+
183
+ return json.dumps({"operation": operation, "a": a, "b": b, "result": result})
184
+
185
+
186
+ # ============================================================================
187
+ # Test Class
188
+ # ============================================================================
189
+
190
+
191
+ @pytest.mark.skipif(SKIP_BEDROCK, reason="AWS credentials not available")
192
+ @pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed")
193
+ class TestHeadroomHookProviderReal:
194
+ """Real-world integration tests for HeadroomHookProvider with Bedrock."""
195
+
196
+ @pytest.fixture
197
+ def bedrock_model(self):
198
+ """Create a BedrockModel instance using Claude 3 Haiku (fast and cheap)."""
199
+ return BedrockModel(
200
+ model_id="anthropic.claude-3-haiku-20240307-v1:0",
201
+ region_name="us-west-2",
202
+ temperature=0.1, # Low temperature for consistent tests
203
+ )
204
+
205
+ @pytest.fixture
206
+ def hook_provider(self):
207
+ """Create a HeadroomHookProvider with test configuration."""
208
+ from headroom.integrations.strands import HeadroomHookProvider
209
+
210
+ return HeadroomHookProvider(
211
+ compress_tool_outputs=True,
212
+ min_tokens_to_compress=50, # Low threshold for testing
213
+ preserve_errors=True,
214
+ )
215
+
216
+ def test_hook_compresses_large_tool_output(self, bedrock_model, hook_provider):
217
+ """Test that large tool outputs are compressed by the hook.
218
+
219
+ This test:
220
+ 1. Creates an agent with the search_logs tool
221
+ 2. Asks a question that triggers the tool
222
+ 3. Verifies the hook compressed the output and saved tokens
223
+ """
224
+ # Create agent with hook provider
225
+ agent = Agent(
226
+ model=bedrock_model,
227
+ tools=[search_logs],
228
+ hooks=[hook_provider],
229
+ )
230
+
231
+ # Ask a question that will trigger the search_logs tool
232
+ result = agent(
233
+ "Search the logs for 'error' and tell me how many entries you found. "
234
+ "Use limit=100 to get plenty of results."
235
+ )
236
+
237
+ # Verify the agent got a response
238
+ assert result is not None
239
+
240
+ # Check hook metrics
241
+ metrics = hook_provider.get_savings_summary()
242
+
243
+ # The hook should have processed at least one tool call
244
+ assert metrics["total_requests"] >= 1, "Hook should have processed tool calls"
245
+
246
+ # With 100 log entries, compression should have occurred
247
+ # and saved significant tokens
248
+ if metrics["compressed_requests"] > 0:
249
+ assert metrics["total_tokens_saved"] > 0, "Should have saved tokens"
250
+ assert metrics["total_tokens_before"] > metrics["total_tokens_after"]
251
+
252
+ def test_hook_preserves_small_outputs(self, bedrock_model, hook_provider):
253
+ """Test that small tool outputs are NOT compressed.
254
+
255
+ This test:
256
+ 1. Creates an agent with a tool returning small output
257
+ 2. Triggers the tool
258
+ 3. Verifies the hook did not modify the small output
259
+ """
260
+ # Reset metrics from any previous tests
261
+ hook_provider.reset()
262
+
263
+ agent = Agent(
264
+ model=bedrock_model,
265
+ tools=[get_small_status],
266
+ hooks=[hook_provider],
267
+ )
268
+
269
+ # Ask a question that will trigger the small status tool
270
+ result = agent("What is the current system status? Use the get_small_status tool.")
271
+
272
+ assert result is not None
273
+
274
+ # Check metrics - small outputs should not be compressed
275
+ metrics = hook_provider.get_savings_summary()
276
+
277
+ # Tool was called but output was below threshold
278
+ if metrics["total_requests"] > 0:
279
+ # For small outputs, tokens_before == tokens_after (no compression)
280
+ for m in hook_provider.metrics_history:
281
+ if m.tool_name == "get_small_status" or "small" in str(m.skip_reason):
282
+ # Either not compressed or skip reason indicates below threshold
283
+ assert not m.was_compressed or m.skip_reason is not None, (
284
+ "Small output should not be compressed"
285
+ )
286
+
287
+ def test_hook_preserves_errors(self, bedrock_model):
288
+ """Test that error results are NOT compressed when preserve_errors=True.
289
+
290
+ This test:
291
+ 1. Creates a hook with preserve_errors=True
292
+ 2. Creates an agent with a tool that returns error data
293
+ 3. Verifies error results are preserved unchanged
294
+ """
295
+ from headroom.integrations.strands import HeadroomHookProvider
296
+
297
+ # Create hook with preserve_errors=True (default)
298
+ hook_with_preserve = HeadroomHookProvider(
299
+ compress_tool_outputs=True,
300
+ min_tokens_to_compress=10, # Very low threshold
301
+ preserve_errors=True,
302
+ )
303
+
304
+ agent = Agent(
305
+ model=bedrock_model,
306
+ tools=[get_error_data],
307
+ hooks=[hook_with_preserve],
308
+ )
309
+
310
+ # Get error data
311
+ result = agent("Get the error data using get_error_data tool and summarize it.")
312
+
313
+ assert result is not None
314
+
315
+ # Check that error-related results were handled appropriately
316
+ metrics = hook_with_preserve.get_savings_summary()
317
+
318
+ # The get_error_data tool returns data about errors but doesn't itself error
319
+ # So it should be processed normally (this tests the flow works)
320
+ assert metrics["total_requests"] >= 0 # May or may not have been called
321
+
322
+ def test_hook_metrics_tracking(self, bedrock_model, hook_provider):
323
+ """Test that metrics are tracked correctly across multiple tool calls.
324
+
325
+ This test:
326
+ 1. Creates an agent with multiple tools
327
+ 2. Makes requests that trigger various tools
328
+ 3. Verifies metrics are accumulated correctly
329
+ """
330
+ # Reset metrics
331
+ hook_provider.reset()
332
+
333
+ agent = Agent(
334
+ model=bedrock_model,
335
+ tools=[search_logs, get_small_status, simple_calculator],
336
+ hooks=[hook_provider],
337
+ )
338
+
339
+ # First request - should trigger search_logs (large output)
340
+ agent("Search logs for 'test' with limit=50 and give me a count.")
341
+
342
+ # Second request - should trigger calculator (small output)
343
+ agent("Calculate 15 + 27 using the calculator tool.")
344
+
345
+ # Third request - should trigger status (small output)
346
+ agent("Get the system status using get_small_status.")
347
+
348
+ # Check accumulated metrics
349
+ metrics = hook_provider.get_savings_summary()
350
+
351
+ # Should have tracked multiple requests
352
+ assert metrics["total_requests"] >= 1, "Should have tracked tool requests"
353
+
354
+ # total_tokens_before should be >= total_tokens_after
355
+ assert metrics["total_tokens_before"] >= metrics["total_tokens_after"]
356
+
357
+ # History should contain records
358
+ history = hook_provider.metrics_history
359
+ assert len(history) >= 1, "Should have metrics history entries"
360
+
361
+ # Each metric should have required fields
362
+ for m in history:
363
+ assert m.request_id is not None
364
+ assert m.timestamp is not None
365
+ assert m.tokens_before >= 0
366
+ assert m.tokens_after >= 0
367
+
368
+ def test_multiple_tool_calls_in_single_request(self, bedrock_model, hook_provider):
369
+ """Test that multiple tool calls in a single agent request are all processed.
370
+
371
+ This test:
372
+ 1. Asks a complex question requiring multiple tools
373
+ 2. Verifies each tool call is processed by the hook
374
+ """
375
+ # Reset metrics
376
+ hook_provider.reset()
377
+
378
+ agent = Agent(
379
+ model=bedrock_model,
380
+ tools=[search_logs, simple_calculator, fetch_user_data],
381
+ hooks=[hook_provider],
382
+ )
383
+
384
+ # Ask a complex question that might trigger multiple tools
385
+ result = agent(
386
+ "I need you to do three things: "
387
+ "1. Search logs for 'api' with limit=30. "
388
+ "2. Calculate 100 * 5 using the calculator. "
389
+ "3. Tell me the total number of results from step 1."
390
+ )
391
+
392
+ assert result is not None
393
+
394
+ # Check that multiple tool calls were processed
395
+ metrics = hook_provider.get_savings_summary()
396
+
397
+ # Should have processed at least the search_logs call
398
+ assert metrics["total_requests"] >= 1
399
+
400
+ # Verify metrics history
401
+ history = hook_provider.metrics_history
402
+
403
+ # At minimum, should have processed search_logs (which has large output)
404
+ # The actual tools called depend on the model's interpretation
405
+ assert len(history) >= 1
406
+
407
+ # Check that we have tool names recorded
408
+ tool_names = [m.tool_name for m in history]
409
+ assert all(name is not None for name in tool_names)
410
+
411
+ def test_hook_reset_clears_metrics(self, bedrock_model, hook_provider):
412
+ """Test that reset() clears all accumulated metrics.
413
+
414
+ This test:
415
+ 1. Makes some requests to accumulate metrics
416
+ 2. Calls reset()
417
+ 3. Verifies all metrics are cleared
418
+ """
419
+ agent = Agent(
420
+ model=bedrock_model,
421
+ tools=[search_logs],
422
+ hooks=[hook_provider],
423
+ )
424
+
425
+ # Make a request to accumulate metrics
426
+ agent("Search logs for 'test' with limit=20.")
427
+
428
+ # Verify we have some metrics
429
+ assert hook_provider.total_tokens_saved >= 0
430
+
431
+ # Reset
432
+ hook_provider.reset()
433
+
434
+ # Verify metrics are cleared
435
+ assert hook_provider.total_tokens_saved == 0
436
+ assert len(hook_provider.metrics_history) == 0
437
+
438
+ metrics = hook_provider.get_savings_summary()
439
+ assert metrics["total_requests"] == 0
440
+ assert metrics["total_tokens_saved"] == 0
441
+
442
+ def test_hook_with_compression_disabled(self, bedrock_model):
443
+ """Test that hook passes through without compression when disabled.
444
+
445
+ This test:
446
+ 1. Creates a hook with compress_tool_outputs=False
447
+ 2. Verifies tool outputs are not modified
448
+ """
449
+ from headroom.integrations.strands import HeadroomHookProvider
450
+
451
+ # Create hook with compression disabled
452
+ disabled_hook = HeadroomHookProvider(
453
+ compress_tool_outputs=False,
454
+ min_tokens_to_compress=10,
455
+ )
456
+
457
+ agent = Agent(
458
+ model=bedrock_model,
459
+ tools=[search_logs],
460
+ hooks=[disabled_hook],
461
+ )
462
+
463
+ result = agent("Search logs for 'api' with limit=50.")
464
+
465
+ assert result is not None
466
+
467
+ # When compression is disabled, no requests should be tracked
468
+ # (the hook doesn't register callbacks when disabled)
469
+ metrics = disabled_hook.get_savings_summary()
470
+ assert metrics["compressed_requests"] == 0
471
+
472
+ def test_hook_concurrent_safety(self, bedrock_model, hook_provider):
473
+ """Test that hook is thread-safe for concurrent access.
474
+
475
+ This test verifies that metrics tracking is thread-safe
476
+ by checking that accumulated values are consistent.
477
+ """
478
+ import threading
479
+
480
+ # Reset metrics
481
+ hook_provider.reset()
482
+
483
+ agent = Agent(
484
+ model=bedrock_model,
485
+ tools=[simple_calculator],
486
+ hooks=[hook_provider],
487
+ )
488
+
489
+ results = []
490
+ errors = []
491
+
492
+ def make_request(n: int):
493
+ try:
494
+ result = agent(f"Calculate {n} + {n} using simple_calculator.")
495
+ results.append(result)
496
+ except Exception as e:
497
+ errors.append(e)
498
+
499
+ # Run a few sequential requests (concurrent Bedrock calls might be rate-limited)
500
+ threads = []
501
+ for i in range(3):
502
+ t = threading.Thread(target=make_request, args=(i,))
503
+ threads.append(t)
504
+ t.start()
505
+ # Small delay to avoid rate limiting
506
+ import time
507
+
508
+ time.sleep(0.5)
509
+
510
+ for t in threads:
511
+ t.join(timeout=60) # 60 second timeout per thread
512
+
513
+ # Check we got results (some may have failed due to rate limits)
514
+ assert len(results) > 0 or len(errors) > 0
515
+
516
+ # Metrics should still be consistent
517
+ metrics = hook_provider.get_savings_summary()
518
+ assert metrics["total_tokens_before"] >= metrics["total_tokens_after"]
519
+
520
+ def test_hook_handles_empty_tool_response(self, bedrock_model, hook_provider):
521
+ """Test that hook handles tools returning empty responses gracefully."""
522
+
523
+ @tool
524
+ def empty_response() -> str:
525
+ """Return an empty response."""
526
+ return ""
527
+
528
+ hook_provider.reset()
529
+
530
+ agent = Agent(
531
+ model=bedrock_model,
532
+ tools=[empty_response],
533
+ hooks=[hook_provider],
534
+ )
535
+
536
+ # This might not trigger the tool if the model decides it's not needed
537
+ result = agent("Call the empty_response tool and tell me what you got.")
538
+
539
+ assert result is not None
540
+
541
+ # Should handle gracefully without errors
542
+ metrics = hook_provider.get_savings_summary()
543
+ # Just verify no exceptions and metrics are valid
544
+ assert metrics["total_tokens_before"] >= 0
545
+ assert metrics["total_tokens_after"] >= 0
tests/integrations/test_strands/test_hooks_unit.py ADDED
@@ -0,0 +1,564 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Unit tests for Strands HeadroomHookProvider.
2
+
3
+ These tests use mocks and do NOT require AWS credentials or strands-agents.
4
+ They test the internal logic of HeadroomHookProvider in isolation.
5
+
6
+ For real integration tests, see test_hooks.py.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import json
12
+ import threading
13
+ from datetime import datetime, timezone
14
+ from unittest.mock import MagicMock
15
+
16
+ import pytest
17
+
18
+ # Check if strands-agents is installed for proper skip handling
19
+ try:
20
+ import strands # noqa: F401
21
+
22
+ STRANDS_AVAILABLE = True
23
+ except ImportError:
24
+ STRANDS_AVAILABLE = False
25
+
26
+
27
+ # Skip all tests if Strands not installed
28
+ pytestmark = pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed")
29
+
30
+
31
+ class TestHeadroomHookProviderInit:
32
+ """Tests for HeadroomHookProvider initialization."""
33
+
34
+ def test_init_with_defaults(self):
35
+ """Initialize with default settings."""
36
+ from headroom.integrations.strands import HeadroomHookProvider
37
+
38
+ hook = HeadroomHookProvider()
39
+
40
+ assert hook.compress_tool_outputs is True
41
+ assert hook.min_tokens_to_compress == 100
42
+ assert hook.preserve_errors is True
43
+ assert hook.total_tokens_saved == 0
44
+ assert hook.metrics_history == []
45
+
46
+ def test_init_with_custom_config(self):
47
+ """Initialize with custom configuration."""
48
+ from headroom import HeadroomConfig
49
+ from headroom.integrations.strands import HeadroomHookProvider
50
+
51
+ config = HeadroomConfig()
52
+ config.smart_crusher.min_tokens_to_crush = 200
53
+ config.smart_crusher.max_items_after_crush = 20
54
+
55
+ hook = HeadroomHookProvider(
56
+ compress_tool_outputs=False,
57
+ min_tokens_to_compress=500,
58
+ config=config,
59
+ preserve_errors=False,
60
+ )
61
+
62
+ assert hook.compress_tool_outputs is False
63
+ assert hook.min_tokens_to_compress == 500
64
+ assert hook.config is config
65
+ assert hook.preserve_errors is False
66
+
67
+ def test_init_creates_default_config_if_none(self):
68
+ """Initialize creates a default HeadroomConfig if none provided."""
69
+ from headroom import HeadroomConfig
70
+ from headroom.integrations.strands import HeadroomHookProvider
71
+
72
+ hook = HeadroomHookProvider()
73
+
74
+ assert hook.config is not None
75
+ assert isinstance(hook.config, HeadroomConfig)
76
+
77
+
78
+ class TestRegisterHooks:
79
+ """Tests for HeadroomHookProvider.register_hooks method."""
80
+
81
+ def test_register_hooks_adds_callback_to_registry(self):
82
+ """register_hooks adds AfterToolCallEvent callback to registry."""
83
+ from headroom.integrations.strands import HeadroomHookProvider
84
+
85
+ hook = HeadroomHookProvider(compress_tool_outputs=True)
86
+ mock_registry = MagicMock()
87
+
88
+ hook.register_hooks(mock_registry)
89
+
90
+ # Should have registered exactly one callback for AfterToolCallEvent
91
+ assert mock_registry.add_callback.call_count == 1
92
+
93
+ def test_register_hooks_skips_when_compression_disabled(self):
94
+ """register_hooks does not register callbacks when compression is disabled."""
95
+ from headroom.integrations.strands import HeadroomHookProvider
96
+
97
+ hook = HeadroomHookProvider(compress_tool_outputs=False)
98
+ mock_registry = MagicMock()
99
+
100
+ hook.register_hooks(mock_registry)
101
+
102
+ # Should not have registered any callbacks
103
+ assert mock_registry.add_callback.call_count == 0
104
+
105
+
106
+ class TestCrusherLazyInit:
107
+ """Tests for SmartCrusher lazy initialization."""
108
+
109
+ def test_crusher_is_lazily_initialized(self):
110
+ """SmartCrusher is not created until first access."""
111
+ from headroom.integrations.strands import HeadroomHookProvider
112
+
113
+ hook = HeadroomHookProvider()
114
+
115
+ # Directly check internal state - crusher should be None initially
116
+ assert hook._crusher is None
117
+
118
+ # Access the crusher property
119
+ crusher = hook.crusher
120
+
121
+ # Now it should be initialized
122
+ assert crusher is not None
123
+ assert hook._crusher is crusher
124
+
125
+ def test_crusher_uses_configured_min_tokens(self):
126
+ """SmartCrusher uses min_tokens_to_compress from hook config."""
127
+ from headroom.integrations.strands import HeadroomHookProvider
128
+
129
+ hook = HeadroomHookProvider(min_tokens_to_compress=250)
130
+
131
+ crusher = hook.crusher
132
+
133
+ # The crusher config should have our min_tokens setting
134
+ assert crusher.config.min_tokens_to_crush == 250
135
+
136
+
137
+ class TestTokenEstimation:
138
+ """Tests for _estimate_tokens helper method."""
139
+
140
+ def test_estimate_tokens_empty_string(self):
141
+ """Estimate returns 0 for empty string."""
142
+ from headroom.integrations.strands import HeadroomHookProvider
143
+
144
+ hook = HeadroomHookProvider()
145
+ assert hook._estimate_tokens("") == 0
146
+
147
+ def test_estimate_tokens_short_string(self):
148
+ """Estimate uses ~4 chars per token heuristic."""
149
+ from headroom.integrations.strands import HeadroomHookProvider
150
+
151
+ hook = HeadroomHookProvider()
152
+
153
+ # 12 chars = 3 tokens (12 // 4)
154
+ assert hook._estimate_tokens("hello world!") == 3
155
+
156
+ # 20 chars = 5 tokens
157
+ assert hook._estimate_tokens("a" * 20) == 5
158
+
159
+
160
+ class TestExtractTextContent:
161
+ """Tests for _extract_text_content helper method."""
162
+
163
+ def test_extract_from_text_content(self):
164
+ """Extract text from content with text field."""
165
+ from headroom.integrations.strands import HeadroomHookProvider
166
+
167
+ hook = HeadroomHookProvider()
168
+ result = {"content": [{"text": "Hello world"}]}
169
+
170
+ extracted = hook._extract_text_content(result)
171
+ assert extracted == "Hello world"
172
+
173
+ def test_extract_from_json_content(self):
174
+ """Extract and serialize JSON content."""
175
+ from headroom.integrations.strands import HeadroomHookProvider
176
+
177
+ hook = HeadroomHookProvider()
178
+ result = {"content": [{"json": {"key": "value"}}]}
179
+
180
+ extracted = hook._extract_text_content(result)
181
+ assert extracted == '{"key": "value"}'
182
+
183
+ def test_extract_empty_content(self):
184
+ """Return empty string for empty content."""
185
+ from headroom.integrations.strands import HeadroomHookProvider
186
+
187
+ hook = HeadroomHookProvider()
188
+ result = {"content": []}
189
+
190
+ extracted = hook._extract_text_content(result)
191
+ assert extracted == ""
192
+
193
+ def test_extract_missing_content(self):
194
+ """Return empty string for missing content key."""
195
+ from headroom.integrations.strands import HeadroomHookProvider
196
+
197
+ hook = HeadroomHookProvider()
198
+ result = {}
199
+
200
+ extracted = hook._extract_text_content(result)
201
+ assert extracted == ""
202
+
203
+
204
+ class TestShouldSkipCompression:
205
+ """Tests for _should_skip_compression helper method."""
206
+
207
+ def test_skip_when_compression_disabled(self):
208
+ """Skip compression when compress_tool_outputs is False."""
209
+ from headroom.integrations.strands import HeadroomHookProvider
210
+
211
+ hook = HeadroomHookProvider(compress_tool_outputs=False)
212
+ result = {"content": [{"text": "data"}]}
213
+
214
+ skip_reason = hook._should_skip_compression(result)
215
+ assert skip_reason == "compression_disabled"
216
+
217
+ def test_skip_error_results_when_preserve_errors_true(self):
218
+ """Skip error results when preserve_errors is True."""
219
+ from headroom.integrations.strands import HeadroomHookProvider
220
+
221
+ hook = HeadroomHookProvider(preserve_errors=True)
222
+ result = {"status": "error", "content": [{"text": "Error message"}]}
223
+
224
+ skip_reason = hook._should_skip_compression(result)
225
+ assert skip_reason == "error_result_preserved"
226
+
227
+ def test_allow_error_results_when_preserve_errors_false(self):
228
+ """Allow error results when preserve_errors is False."""
229
+ from headroom.integrations.strands import HeadroomHookProvider
230
+
231
+ hook = HeadroomHookProvider(preserve_errors=False)
232
+ result = {"status": "error", "content": [{"text": "Error message"}]}
233
+
234
+ skip_reason = hook._should_skip_compression(result)
235
+ assert skip_reason is None
236
+
237
+ def test_skip_empty_content(self):
238
+ """Skip results with empty content."""
239
+ from headroom.integrations.strands import HeadroomHookProvider
240
+
241
+ hook = HeadroomHookProvider()
242
+ result = {"content": []}
243
+
244
+ skip_reason = hook._should_skip_compression(result)
245
+ assert skip_reason == "empty_content"
246
+
247
+ def test_allow_valid_content(self):
248
+ """Allow results with valid content."""
249
+ from headroom.integrations.strands import HeadroomHookProvider
250
+
251
+ hook = HeadroomHookProvider()
252
+ result = {"content": [{"text": "some data"}]}
253
+
254
+ skip_reason = hook._should_skip_compression(result)
255
+ assert skip_reason is None
256
+
257
+
258
+ class TestCompressToolResult:
259
+ """Tests for _compress_tool_result hook handler."""
260
+
261
+ def test_compress_large_tool_output(self):
262
+ """Compresses large tool output and tracks metrics."""
263
+ from headroom.integrations.strands import HeadroomHookProvider
264
+
265
+ hook = HeadroomHookProvider(
266
+ compress_tool_outputs=True,
267
+ min_tokens_to_compress=10, # Low threshold for testing
268
+ )
269
+
270
+ # Create large JSON output (50 items)
271
+ large_data = [{"id": i, "value": f"item-{i}", "data": "x" * 50} for i in range(50)]
272
+ large_json = json.dumps(large_data)
273
+
274
+ mock_event = MagicMock()
275
+ mock_event.tool_use = {"name": "get_items", "toolUseId": "tool-123"}
276
+ mock_event.result = {"content": [{"text": large_json}]}
277
+
278
+ hook._compress_tool_result(mock_event)
279
+
280
+ # Verify metrics were recorded
281
+ assert len(hook.metrics_history) == 1
282
+ metrics = hook.metrics_history[0]
283
+ assert metrics.tool_name == "get_items"
284
+ assert metrics.tool_use_id == "tool-123"
285
+ assert metrics.tokens_before > 0
286
+
287
+ def test_skip_compression_below_threshold(self):
288
+ """Does not compress output below token threshold."""
289
+ from headroom.integrations.strands import HeadroomHookProvider
290
+
291
+ hook = HeadroomHookProvider(
292
+ compress_tool_outputs=True,
293
+ min_tokens_to_compress=10000, # High threshold
294
+ )
295
+
296
+ mock_event = MagicMock()
297
+ mock_event.tool_use = {"name": "small_tool", "toolUseId": "tool-456"}
298
+ mock_event.result = {"content": [{"text": '{"status": "ok"}'}]}
299
+
300
+ hook._compress_tool_result(mock_event)
301
+
302
+ # Metrics should show skipped compression
303
+ assert len(hook.metrics_history) == 1
304
+ metrics = hook.metrics_history[0]
305
+ assert metrics.was_compressed is False
306
+ assert "below_threshold" in metrics.skip_reason
307
+
308
+ def test_skip_compression_when_disabled(self):
309
+ """Does not compress when compression is disabled."""
310
+ from headroom.integrations.strands import HeadroomHookProvider
311
+
312
+ hook = HeadroomHookProvider(compress_tool_outputs=False)
313
+
314
+ mock_event = MagicMock()
315
+ mock_event.tool_use = {"name": "test_tool", "toolUseId": "tool-789"}
316
+ mock_event.result = {"content": [{"text": '{"data": "value"}'}]}
317
+
318
+ hook._compress_tool_result(mock_event)
319
+
320
+ # Metrics should show compression disabled
321
+ assert len(hook.metrics_history) == 1
322
+ metrics = hook.metrics_history[0]
323
+ assert metrics.was_compressed is False
324
+ assert metrics.skip_reason == "compression_disabled"
325
+
326
+
327
+ class TestMetricsTracking:
328
+ """Tests for metrics tracking and aggregation."""
329
+
330
+ def test_total_tokens_saved_accumulates(self):
331
+ """total_tokens_saved accumulates across compressions."""
332
+ from headroom.integrations.strands import HeadroomHookProvider
333
+
334
+ hook = HeadroomHookProvider(
335
+ compress_tool_outputs=True,
336
+ min_tokens_to_compress=10,
337
+ )
338
+
339
+ # Simulate two compressions with savings
340
+ for i in range(2):
341
+ large_data = [{"id": j, "data": "x" * 100} for j in range(50)]
342
+ mock_event = MagicMock()
343
+ mock_event.tool_use = {"name": f"tool_{i}", "toolUseId": f"id_{i}"}
344
+ mock_event.result = {"content": [{"text": json.dumps(large_data)}]}
345
+
346
+ hook._compress_tool_result(mock_event)
347
+
348
+ # Should have accumulated some savings
349
+ compressed_count = sum(1 for m in hook.metrics_history if m.was_compressed)
350
+ if compressed_count > 0:
351
+ assert hook.total_tokens_saved >= 0
352
+
353
+ def test_metrics_history_bounded_to_100(self):
354
+ """metrics_history keeps only last 100 entries."""
355
+ from headroom.integrations.strands import HeadroomHookProvider
356
+
357
+ hook = HeadroomHookProvider(
358
+ compress_tool_outputs=True,
359
+ min_tokens_to_compress=10,
360
+ )
361
+
362
+ # Directly add 150 metrics
363
+ for i in range(150):
364
+ hook._record_metrics(
365
+ request_id=f"req_{i}",
366
+ tool_name=f"tool_{i}",
367
+ tool_use_id=f"id_{i}",
368
+ tokens_before=100,
369
+ tokens_after=50,
370
+ was_compressed=True,
371
+ skip_reason=None,
372
+ )
373
+
374
+ # Should be bounded at 100
375
+ assert len(hook.metrics_history) == 100
376
+
377
+ # Should contain the most recent entries
378
+ last_metric = hook.metrics_history[-1]
379
+ assert last_metric.request_id == "req_149"
380
+
381
+
382
+ class TestGetSavingsSummary:
383
+ """Tests for get_savings_summary method."""
384
+
385
+ def test_empty_summary(self):
386
+ """Returns zero values when no metrics recorded."""
387
+ from headroom.integrations.strands import HeadroomHookProvider
388
+
389
+ hook = HeadroomHookProvider()
390
+ summary = hook.get_savings_summary()
391
+
392
+ assert summary["total_requests"] == 0
393
+ assert summary["compressed_requests"] == 0
394
+ assert summary["total_tokens_saved"] == 0
395
+ assert summary["average_savings_percent"] == 0.0
396
+
397
+ def test_summary_with_compressions(self):
398
+ """Returns correct summary with recorded compressions."""
399
+ from headroom.integrations.strands import HeadroomHookProvider
400
+ from headroom.integrations.strands.hooks import CompressionMetrics
401
+
402
+ hook = HeadroomHookProvider()
403
+
404
+ # Add metrics manually
405
+ hook._metrics_history = [
406
+ CompressionMetrics(
407
+ request_id="1",
408
+ timestamp=datetime.now(timezone.utc),
409
+ tool_name="tool_a",
410
+ tool_use_id="id_1",
411
+ tokens_before=100,
412
+ tokens_after=60,
413
+ tokens_saved=40,
414
+ savings_percent=40.0,
415
+ was_compressed=True,
416
+ skip_reason=None,
417
+ ),
418
+ CompressionMetrics(
419
+ request_id="2",
420
+ timestamp=datetime.now(timezone.utc),
421
+ tool_name="tool_b",
422
+ tool_use_id="id_2",
423
+ tokens_before=200,
424
+ tokens_after=100,
425
+ tokens_saved=100,
426
+ savings_percent=50.0,
427
+ was_compressed=True,
428
+ skip_reason=None,
429
+ ),
430
+ CompressionMetrics(
431
+ request_id="3",
432
+ timestamp=datetime.now(timezone.utc),
433
+ tool_name="tool_c",
434
+ tool_use_id="id_3",
435
+ tokens_before=50,
436
+ tokens_after=50,
437
+ tokens_saved=0,
438
+ savings_percent=0.0,
439
+ was_compressed=False,
440
+ skip_reason="below_threshold",
441
+ ),
442
+ ]
443
+ hook._total_tokens_saved = 140
444
+
445
+ summary = hook.get_savings_summary()
446
+
447
+ assert summary["total_requests"] == 3
448
+ assert summary["compressed_requests"] == 2
449
+ assert summary["total_tokens_saved"] == 140
450
+ assert summary["average_savings_percent"] == 45.0 # (40 + 50) / 2
451
+ assert summary["total_tokens_before"] == 350
452
+ assert summary["total_tokens_after"] == 210
453
+
454
+
455
+ class TestReset:
456
+ """Tests for reset method."""
457
+
458
+ def test_reset_clears_all_state(self):
459
+ """reset() clears all tracked state."""
460
+ from headroom.integrations.strands import HeadroomHookProvider
461
+ from headroom.integrations.strands.hooks import CompressionMetrics
462
+
463
+ hook = HeadroomHookProvider()
464
+
465
+ # Add some state
466
+ hook._metrics_history = [
467
+ CompressionMetrics(
468
+ request_id="1",
469
+ timestamp=datetime.now(timezone.utc),
470
+ tool_name="test",
471
+ tool_use_id="id_1",
472
+ tokens_before=100,
473
+ tokens_after=50,
474
+ tokens_saved=50,
475
+ savings_percent=50.0,
476
+ was_compressed=True,
477
+ )
478
+ ]
479
+ hook._total_tokens_saved = 50
480
+
481
+ # Reset
482
+ hook.reset()
483
+
484
+ # Verify all state cleared
485
+ assert hook._metrics_history == []
486
+ assert hook._total_tokens_saved == 0
487
+ assert hook.total_tokens_saved == 0
488
+ assert len(hook.metrics_history) == 0
489
+
490
+
491
+ class TestThreadSafety:
492
+ """Tests for thread-safety of metrics tracking."""
493
+
494
+ def test_concurrent_metric_recording(self):
495
+ """Metrics recording is thread-safe."""
496
+ from headroom.integrations.strands import HeadroomHookProvider
497
+
498
+ hook = HeadroomHookProvider()
499
+
500
+ def record_metrics(thread_id):
501
+ for i in range(10):
502
+ hook._record_metrics(
503
+ request_id=f"thread_{thread_id}_req_{i}",
504
+ tool_name=f"tool_{thread_id}_{i}",
505
+ tool_use_id=f"id_{thread_id}_{i}",
506
+ tokens_before=100,
507
+ tokens_after=50,
508
+ was_compressed=True,
509
+ skip_reason=None,
510
+ )
511
+
512
+ threads = []
513
+ for t_id in range(5):
514
+ t = threading.Thread(target=record_metrics, args=(t_id,))
515
+ threads.append(t)
516
+ t.start()
517
+
518
+ for t in threads:
519
+ t.join()
520
+
521
+ # Should have recorded 50 metrics (5 threads * 10 each)
522
+ # But bounded to 100, so if we had more it would be truncated
523
+ assert len(hook.metrics_history) == 50
524
+ assert hook.total_tokens_saved == 50 * 50 # 50 metrics * 50 tokens each
525
+
526
+
527
+ class TestUpdateResultContent:
528
+ """Tests for _update_result_content helper method."""
529
+
530
+ def test_update_preserves_json_structure(self):
531
+ """Updates preserve JSON structure when possible."""
532
+ from headroom.integrations.strands import HeadroomHookProvider
533
+
534
+ hook = HeadroomHookProvider()
535
+ result = {"content": [{"json": {"original": "data"}}]}
536
+
537
+ compressed = '{"compressed": "data"}'
538
+ hook._update_result_content(result, compressed)
539
+
540
+ # Should update with parsed JSON
541
+ assert result["content"] == [{"json": {"compressed": "data"}}]
542
+
543
+ def test_update_uses_text_for_non_json(self):
544
+ """Updates use text format for non-JSON content."""
545
+ from headroom.integrations.strands import HeadroomHookProvider
546
+
547
+ hook = HeadroomHookProvider()
548
+ result = {"content": [{"text": "original text"}]}
549
+
550
+ compressed = "compressed text"
551
+ hook._update_result_content(result, compressed)
552
+
553
+ assert result["content"] == [{"text": "compressed text"}]
554
+
555
+ def test_update_creates_content_if_empty(self):
556
+ """Creates content list if missing."""
557
+ from headroom.integrations.strands import HeadroomHookProvider
558
+
559
+ hook = HeadroomHookProvider()
560
+ result = {"content": []}
561
+
562
+ hook._update_result_content(result, "new content")
563
+
564
+ assert result["content"] == [{"text": "new content"}]
tests/integrations/test_strands/test_model.py ADDED
@@ -0,0 +1,673 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Real-world integration tests for Strands HeadroomStrandsModel.
2
+
3
+ These tests use actual AWS Bedrock API calls with real credentials.
4
+ NO MOCKS - all tests hit the real Bedrock API.
5
+
6
+ Skip in CI if AWS credentials are not available.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import json
12
+ import os
13
+
14
+ import pytest
15
+
16
+ # Check for AWS credentials availability
17
+ SKIP_BEDROCK = not (
18
+ os.environ.get("AWS_ACCESS_KEY_ID")
19
+ or os.environ.get("AWS_PROFILE")
20
+ or os.path.exists(os.path.expanduser("~/.aws/credentials"))
21
+ )
22
+
23
+ # Check if strands-agents is installed
24
+ try:
25
+ from strands import Agent, tool
26
+ from strands.models import BedrockModel
27
+
28
+ STRANDS_AVAILABLE = True
29
+ except ImportError:
30
+ STRANDS_AVAILABLE = False
31
+
32
+ # Provide a no-op decorator when strands is not installed
33
+ def tool(fn):
34
+ return fn
35
+
36
+ Agent = None # type: ignore
37
+ BedrockModel = None # type: ignore
38
+
39
+ # Skip all tests if dependencies not available
40
+ pytestmark = [
41
+ pytest.mark.skipif(SKIP_BEDROCK, reason="AWS credentials not available"),
42
+ pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed"),
43
+ ]
44
+
45
+
46
+ # ============================================================================
47
+ # Test Tools - Generate realistic data for optimization testing
48
+ # These are defined with @tool decorator for use when strands is installed.
49
+ # When strands is not installed, the no-op decorator ensures import succeeds.
50
+ # ============================================================================
51
+
52
+
53
+ @tool
54
+ def get_database_records(table: str, limit: int = 50) -> str:
55
+ """Fetch records from a database table. Returns JSON array.
56
+
57
+ Args:
58
+ table: Name of the database table
59
+ limit: Maximum records to return
60
+
61
+ Returns:
62
+ JSON array of database records
63
+ """
64
+ records = [
65
+ {
66
+ "id": i,
67
+ "table": table,
68
+ "created_at": f"2024-01-{(i % 28) + 1:02d}T{10 + (i % 12):02d}:00:00Z",
69
+ "updated_at": f"2024-01-{(i % 28) + 1:02d}T{11 + (i % 12):02d}:00:00Z",
70
+ "status": ["active", "inactive", "pending", "archived"][i % 4],
71
+ "priority": ["low", "medium", "high", "critical"][i % 4],
72
+ "data": {
73
+ "field1": f"value_{i}_{table}",
74
+ "field2": i * 100,
75
+ "field3": i % 2 == 0,
76
+ "metadata": {
77
+ "source": "database",
78
+ "version": f"1.{i % 10}.0",
79
+ "tags": [f"tag_{j}" for j in range(i % 5 + 1)],
80
+ },
81
+ },
82
+ "metrics": {
83
+ "read_count": i * 10,
84
+ "write_count": i * 5,
85
+ "error_count": i % 3,
86
+ "latency_ms": 50 + (i * 7) % 200,
87
+ },
88
+ }
89
+ for i in range(limit)
90
+ ]
91
+ return json.dumps(records, indent=2)
92
+
93
+
94
+ @tool
95
+ def get_large_logs(query: str, count: int = 200) -> str:
96
+ """Fetch verbose log data that should trigger compression.
97
+
98
+ Args:
99
+ query: Search query for logs
100
+ count: Number of log entries to return
101
+
102
+ Returns:
103
+ JSON array of detailed log entries
104
+ """
105
+ logs = [
106
+ {
107
+ "log_id": f"log_{i:08d}",
108
+ "timestamp": f"2024-01-{(i % 28) + 1:02d}T{10 + (i % 12):02d}:{i % 60:02d}:00Z",
109
+ "level": ["DEBUG", "INFO", "WARN", "ERROR"][i % 4],
110
+ "service": f"service_{i % 10}",
111
+ "message": f"Processing request for query '{query}' - step {i}",
112
+ "request_id": f"req_{i:012d}",
113
+ "trace_id": f"trace_{i:016x}",
114
+ "span_id": f"span_{i:08x}",
115
+ "user_id": f"user_{i % 100:04d}",
116
+ "session_id": f"sess_{i:010d}",
117
+ "metadata": {
118
+ "host": f"server-{i % 20:02d}.example.com",
119
+ "region": ["us-west-2", "us-east-1", "eu-west-1", "ap-southeast-1"][i % 4],
120
+ "instance_type": ["t3.micro", "t3.small", "t3.medium", "t3.large"][i % 4],
121
+ "container_id": f"container_{i:08x}",
122
+ "kubernetes_pod": f"pod-{i:06d}",
123
+ "kubernetes_namespace": "production",
124
+ },
125
+ "metrics": {
126
+ "duration_ms": 50 + (i * 3) % 500,
127
+ "memory_mb": 128 + (i * 7) % 1024,
128
+ "cpu_percent": 5 + (i * 2) % 95,
129
+ "network_bytes_in": i * 1024,
130
+ "network_bytes_out": i * 512,
131
+ },
132
+ "tags": ["env:prod", f"version:1.{i % 10}.0", "team:backend"],
133
+ }
134
+ for i in range(count)
135
+ ]
136
+ return json.dumps(logs, indent=2)
137
+
138
+
139
+ @tool
140
+ def analyze_metrics(metric_type: str) -> str:
141
+ """Analyze system metrics. Returns detailed metrics data.
142
+
143
+ Args:
144
+ metric_type: Type of metrics to analyze (cpu, memory, network, disk)
145
+
146
+ Returns:
147
+ JSON object with metric analysis
148
+ """
149
+ data_points = [
150
+ {
151
+ "timestamp": f"2024-01-15T{10 + (i % 12):02d}:{(i * 5) % 60:02d}:00Z",
152
+ "value": 20 + (i * 3) % 80,
153
+ "unit": {"cpu": "%", "memory": "MB", "network": "Mbps", "disk": "GB"}.get(
154
+ metric_type, "units"
155
+ ),
156
+ "host": f"server-{(i % 5) + 1:02d}",
157
+ "region": ["us-west-2", "us-east-1", "eu-west-1"][i % 3],
158
+ "metadata": {
159
+ "collection_interval": 60,
160
+ "aggregation": "avg",
161
+ "quality": "good" if i % 5 != 0 else "degraded",
162
+ },
163
+ }
164
+ for i in range(100)
165
+ ]
166
+
167
+ return json.dumps(
168
+ {
169
+ "metric_type": metric_type,
170
+ "time_range": {"start": "2024-01-15T10:00:00Z", "end": "2024-01-15T22:00:00Z"},
171
+ "data_points": data_points,
172
+ "summary": {
173
+ "min": 20,
174
+ "max": 99,
175
+ "avg": 55.5,
176
+ "p50": 52,
177
+ "p95": 90,
178
+ "p99": 97,
179
+ },
180
+ },
181
+ indent=2,
182
+ )
183
+
184
+
185
+ @tool
186
+ def quick_lookup(key: str) -> str:
187
+ """Quick key-value lookup. Returns small response.
188
+
189
+ Args:
190
+ key: The key to look up
191
+
192
+ Returns:
193
+ Small JSON with the value
194
+ """
195
+ return json.dumps({"key": key, "value": f"result_for_{key}", "found": True})
196
+
197
+
198
+ @tool
199
+ def math_operation(x: float, y: float, op: str) -> str:
200
+ """Perform a math operation.
201
+
202
+ Args:
203
+ x: First operand
204
+ y: Second operand
205
+ op: Operation (add, sub, mul, div)
206
+
207
+ Returns:
208
+ Result of the operation
209
+ """
210
+ operations = {
211
+ "add": x + y,
212
+ "sub": x - y,
213
+ "mul": x * y,
214
+ "div": x / y if y != 0 else None,
215
+ }
216
+ result = operations.get(op, None)
217
+ return json.dumps({"x": x, "y": y, "operation": op, "result": result})
218
+
219
+
220
+ # ============================================================================
221
+ # Test Class for HeadroomStrandsModel
222
+ # ============================================================================
223
+
224
+
225
+ @pytest.mark.skipif(SKIP_BEDROCK, reason="AWS credentials not available")
226
+ @pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed")
227
+ class TestHeadroomStrandsModelReal:
228
+ """Real-world integration tests for HeadroomStrandsModel with Bedrock."""
229
+
230
+ @pytest.fixture
231
+ def base_bedrock_model(self):
232
+ """Create a base BedrockModel instance using Claude 3 Haiku (fast and cheap)."""
233
+ return BedrockModel(
234
+ model_id="anthropic.claude-3-haiku-20240307-v1:0",
235
+ region_name="us-west-2",
236
+ temperature=0.1,
237
+ )
238
+
239
+ @pytest.fixture
240
+ def wrapped_model(self, base_bedrock_model):
241
+ """Create a HeadroomStrandsModel wrapping the Bedrock model."""
242
+ from headroom.integrations.strands import HeadroomStrandsModel
243
+
244
+ return HeadroomStrandsModel(
245
+ wrapped_model=base_bedrock_model,
246
+ auto_detect_provider=True,
247
+ )
248
+
249
+ def test_stream_returns_proper_events(self, wrapped_model):
250
+ """Test that stream() works and returns proper StreamEvents.
251
+
252
+ The Strands Agent uses the model's stream() method internally.
253
+ This test verifies that the wrapped model properly streams responses.
254
+ """
255
+ wrapped_model.reset()
256
+
257
+ agent = Agent(model=wrapped_model)
258
+
259
+ # Make a request - the agent internally calls stream() on the model
260
+ result = agent("Count from 1 to 5, one number per line.")
261
+
262
+ # Verify we got a response (proves streaming worked)
263
+ assert result is not None
264
+ response_text = str(result)
265
+ assert len(response_text) > 0
266
+
267
+ # The response should contain numbers 1-5
268
+ for num in ["1", "2", "3", "4", "5"]:
269
+ assert num in response_text, f"Expected {num} in response"
270
+
271
+ # Metrics should be tracked (proves stream() was intercepted properly)
272
+ metrics = wrapped_model.get_savings_summary()
273
+ assert metrics["total_requests"] >= 1, "stream() should track requests"
274
+
275
+ def test_messages_optimized_large_conversations(self, wrapped_model):
276
+ """Test that messages are actually optimized (tokens_before > tokens_after for large conversations).
277
+
278
+ This test builds up a large conversation context through tool calls
279
+ with verbose JSON responses, then verifies that optimization occurs.
280
+ """
281
+ wrapped_model.reset()
282
+
283
+ agent = Agent(model=wrapped_model, tools=[get_large_logs, get_database_records])
284
+
285
+ # First request - get large logs (200 entries with verbose data)
286
+ agent(
287
+ "Search for logs containing 'error' and get 200 entries using get_large_logs. "
288
+ "Tell me how many ERROR level logs there are."
289
+ )
290
+
291
+ # Second request - more tool output, context grows
292
+ agent(
293
+ "Now get 100 records from the 'events' table using get_database_records. "
294
+ "How many records have 'active' status?"
295
+ )
296
+
297
+ # Third request - even more context
298
+ agent(
299
+ "Based on all the data you've seen, give me a one-sentence summary "
300
+ "of the system health."
301
+ )
302
+
303
+ # Check optimization metrics
304
+ metrics = wrapped_model.get_savings_summary()
305
+
306
+ # Should have processed multiple requests
307
+ assert metrics["total_requests"] >= 1, "Should have processed requests"
308
+
309
+ # With large tool outputs, tokens_before should be significant
310
+ assert metrics["total_tokens_before"] > 0, "Should have counted input tokens"
311
+
312
+ # The key assertion: optimization should reduce tokens
313
+ # (tokens_before >= tokens_after, with strict > when there's compressible content)
314
+ assert metrics["total_tokens_before"] >= metrics["total_tokens_after"], (
315
+ f"Optimization should not increase tokens: "
316
+ f"before={metrics['total_tokens_before']}, after={metrics['total_tokens_after']}"
317
+ )
318
+
319
+ # Check history shows optimization was tracked
320
+ history = wrapped_model.metrics_history
321
+ assert len(history) >= 1, "Should have metrics history"
322
+
323
+ # Verify individual requests track before/after properly
324
+ for m in history:
325
+ assert m.tokens_before >= m.tokens_after, (
326
+ f"Each request should have tokens_before >= tokens_after: "
327
+ f"request_id={m.request_id}, before={m.tokens_before}, after={m.tokens_after}"
328
+ )
329
+
330
+ def test_get_savings_summary_returns_correct_metrics(self, wrapped_model):
331
+ """Test that get_savings_summary() returns correct metrics.
332
+
333
+ Verifies the structure and accuracy of the savings summary.
334
+ """
335
+ wrapped_model.reset()
336
+
337
+ agent = Agent(model=wrapped_model, tools=[get_database_records])
338
+
339
+ # Make a few requests
340
+ agent("Get 30 records from 'users' table.")
341
+ agent("Get 30 records from 'orders' table.")
342
+
343
+ # Get the summary
344
+ summary = wrapped_model.get_savings_summary()
345
+
346
+ # Verify required keys exist
347
+ required_keys = [
348
+ "total_requests",
349
+ "total_tokens_saved",
350
+ "average_savings_percent",
351
+ "total_tokens_before",
352
+ "total_tokens_after",
353
+ ]
354
+ for key in required_keys:
355
+ assert key in summary, f"Summary missing required key: {key}"
356
+
357
+ # Verify values are sensible
358
+ assert summary["total_requests"] >= 1, "Should have at least one request"
359
+ assert summary["total_tokens_before"] >= 0, "tokens_before should be non-negative"
360
+ assert summary["total_tokens_after"] >= 0, "tokens_after should be non-negative"
361
+ assert summary["total_tokens_saved"] >= 0, "tokens_saved should be non-negative"
362
+ assert 0 <= summary["average_savings_percent"] <= 100, (
363
+ "average_savings_percent should be between 0 and 100"
364
+ )
365
+
366
+ # Verify mathematical consistency
367
+ expected_saved = summary["total_tokens_before"] - summary["total_tokens_after"]
368
+ assert summary["total_tokens_saved"] == expected_saved, (
369
+ f"tokens_saved should equal tokens_before - tokens_after: "
370
+ f"saved={summary['total_tokens_saved']}, expected={expected_saved}"
371
+ )
372
+
373
+ def test_reset_clears_all_metrics(self, wrapped_model):
374
+ """Test that reset() clears all accumulated metrics.
375
+
376
+ Verifies that reset() properly clears:
377
+ - total_tokens_saved
378
+ - metrics_history
379
+ - The summary returned by get_savings_summary()
380
+ """
381
+ # Make some requests to accumulate metrics
382
+ agent = Agent(model=wrapped_model)
383
+ agent("Say 'hello world'")
384
+ agent("Say 'goodbye world'")
385
+
386
+ # Verify we have metrics before reset
387
+ assert wrapped_model.total_tokens_saved >= 0
388
+ pre_reset_requests = wrapped_model.get_savings_summary()["total_requests"]
389
+ assert pre_reset_requests >= 1, "Should have requests before reset"
390
+
391
+ # Call reset
392
+ wrapped_model.reset()
393
+
394
+ # Verify all metrics are cleared
395
+ assert wrapped_model.total_tokens_saved == 0, "total_tokens_saved should be 0 after reset"
396
+ assert len(wrapped_model.metrics_history) == 0, (
397
+ "metrics_history should be empty after reset"
398
+ )
399
+
400
+ # Verify get_savings_summary reflects the reset
401
+ summary = wrapped_model.get_savings_summary()
402
+ assert summary["total_requests"] == 0, "total_requests should be 0 after reset"
403
+ assert summary["total_tokens_saved"] == 0, "total_tokens_saved should be 0 after reset"
404
+ assert summary["total_tokens_before"] == 0, "total_tokens_before should be 0 after reset"
405
+ assert summary["total_tokens_after"] == 0, "total_tokens_after should be 0 after reset"
406
+
407
+ # Verify we can still make requests after reset
408
+ agent = Agent(model=wrapped_model)
409
+ agent("Say 'post-reset test'")
410
+
411
+ post_reset_summary = wrapped_model.get_savings_summary()
412
+ assert post_reset_summary["total_requests"] >= 1, "Should track requests after reset"
413
+
414
+ def test_model_wrapper_basic_response(self, wrapped_model):
415
+ """Test that wrapped model produces valid responses."""
416
+ agent = Agent(model=wrapped_model)
417
+
418
+ result = agent("Say 'Hello, Headroom!' and nothing else.")
419
+
420
+ assert result is not None
421
+ content = str(result)
422
+ assert len(content) > 0
423
+
424
+ def test_model_wrapper_with_tools(self, wrapped_model):
425
+ """Test that wrapped model works correctly with tools."""
426
+ wrapped_model.reset()
427
+
428
+ agent = Agent(model=wrapped_model, tools=[quick_lookup, math_operation, analyze_metrics])
429
+
430
+ result = agent(
431
+ "Please do these tasks: "
432
+ "1. Look up the key 'config_setting' using quick_lookup. "
433
+ "2. Calculate 15.5 multiplied by 4 using math_operation. "
434
+ "3. Tell me the results."
435
+ )
436
+
437
+ assert result is not None
438
+
439
+ metrics = wrapped_model.get_savings_summary()
440
+ assert metrics["total_requests"] >= 1
441
+
442
+ def test_model_wrapper_metrics_tracking(self, wrapped_model):
443
+ """Test that metrics are accurately tracked across requests."""
444
+ wrapped_model.reset()
445
+
446
+ agent = Agent(model=wrapped_model, tools=[get_database_records])
447
+
448
+ # Make several requests
449
+ agent("Get 20 records from 'products' table.")
450
+ agent("Get 20 records from 'customers' table.")
451
+ agent("Summarize both sets of records.")
452
+
453
+ metrics = wrapped_model.get_savings_summary()
454
+
455
+ assert metrics["total_requests"] >= 1
456
+ assert metrics["total_tokens_before"] >= metrics["total_tokens_after"]
457
+
458
+ if metrics["total_tokens_saved"] > 0:
459
+ assert metrics["average_savings_percent"] >= 0
460
+ assert metrics["average_savings_percent"] <= 100
461
+
462
+ # History should be bounded
463
+ assert len(wrapped_model.metrics_history) <= 100
464
+
465
+ def test_model_wrapper_attribute_forwarding(self, base_bedrock_model):
466
+ """Test that attributes are forwarded to wrapped model."""
467
+ from headroom.integrations.strands import HeadroomStrandsModel
468
+
469
+ wrapped = HeadroomStrandsModel(
470
+ wrapped_model=base_bedrock_model,
471
+ auto_detect_provider=True,
472
+ )
473
+
474
+ # The wrapper should forward config to the wrapped model (Strands stores model_id in config)
475
+ assert hasattr(wrapped, "config")
476
+ config = wrapped.config
477
+ assert isinstance(config, dict)
478
+ assert "model_id" in config
479
+
480
+ # Access wrapped model directly
481
+ assert wrapped.wrapped_model is base_bedrock_model
482
+
483
+ def test_model_wrapper_custom_config(self, base_bedrock_model):
484
+ """Test that custom HeadroomConfig is applied."""
485
+ from headroom import HeadroomConfig
486
+ from headroom.integrations.strands import HeadroomStrandsModel
487
+
488
+ custom_config = HeadroomConfig()
489
+ custom_config.smart_crusher.min_tokens_to_crush = 50
490
+ custom_config.smart_crusher.max_items_after_crush = 10
491
+
492
+ wrapped = HeadroomStrandsModel(
493
+ wrapped_model=base_bedrock_model,
494
+ config=custom_config,
495
+ auto_detect_provider=True,
496
+ )
497
+
498
+ assert wrapped.headroom_config is custom_config
499
+ assert wrapped.headroom_config.smart_crusher.min_tokens_to_crush == 50
500
+
501
+ # The model should still work
502
+ agent = Agent(model=wrapped)
503
+ result = agent("Say 'test'")
504
+ assert result is not None
505
+
506
+ def test_model_wrapper_provider_detection(self, base_bedrock_model):
507
+ """Test that provider is auto-detected correctly for Bedrock Claude."""
508
+ from headroom.integrations.strands import HeadroomStrandsModel
509
+ from headroom.providers import AnthropicProvider
510
+
511
+ wrapped = HeadroomStrandsModel(
512
+ wrapped_model=base_bedrock_model,
513
+ auto_detect_provider=True,
514
+ )
515
+
516
+ # Access pipeline to trigger lazy initialization
517
+ _ = wrapped.pipeline
518
+
519
+ # For Bedrock Claude models, should detect Anthropic provider
520
+ assert wrapped._headroom_provider is not None
521
+ assert isinstance(wrapped._headroom_provider, AnthropicProvider)
522
+
523
+ def test_model_wrapper_handles_large_context(self, wrapped_model):
524
+ """Test that wrapper handles large context appropriately."""
525
+ wrapped_model.reset()
526
+
527
+ agent = Agent(model=wrapped_model, tools=[analyze_metrics, get_database_records])
528
+
529
+ # Build up context with large tool outputs
530
+ agent("Analyze CPU metrics using analyze_metrics.")
531
+ agent("Get 50 records from 'logs' table using get_database_records.")
532
+ agent("Based on everything, what patterns do you see?")
533
+
534
+ metrics = wrapped_model.get_savings_summary()
535
+ assert metrics["total_requests"] >= 1
536
+ assert metrics["total_tokens_before"] > 0
537
+
538
+ def test_model_wrapper_empty_messages(self, base_bedrock_model):
539
+ """Test that wrapper handles edge cases gracefully."""
540
+ from headroom.integrations.strands import HeadroomStrandsModel
541
+
542
+ wrapped = HeadroomStrandsModel(
543
+ wrapped_model=base_bedrock_model,
544
+ auto_detect_provider=True,
545
+ )
546
+
547
+ # Test with minimal input
548
+ agent = Agent(model=wrapped)
549
+ result = agent("Hi")
550
+
551
+ assert result is not None
552
+
553
+ def test_model_wrapper_thread_safety(self, base_bedrock_model):
554
+ """Test that wrapper is thread-safe for metrics tracking."""
555
+ import threading
556
+ import time
557
+
558
+ from headroom.integrations.strands import HeadroomStrandsModel
559
+
560
+ wrapped = HeadroomStrandsModel(
561
+ wrapped_model=base_bedrock_model,
562
+ auto_detect_provider=True,
563
+ )
564
+
565
+ agent = Agent(model=wrapped)
566
+
567
+ results = []
568
+ errors = []
569
+
570
+ def make_request(msg: str):
571
+ try:
572
+ result = agent(msg)
573
+ results.append(result)
574
+ except Exception as e:
575
+ errors.append(e)
576
+
577
+ threads = []
578
+ messages = ["Say 'one'", "Say 'two'", "Say 'three'"]
579
+
580
+ for msg in messages:
581
+ t = threading.Thread(target=make_request, args=(msg,))
582
+ threads.append(t)
583
+ t.start()
584
+ time.sleep(0.5) # Small delay to avoid rate limiting
585
+
586
+ for t in threads:
587
+ t.join(timeout=60)
588
+
589
+ # Should have some results (may have errors due to rate limiting)
590
+ assert len(results) > 0 or len(errors) > 0
591
+
592
+ # Metrics should be consistent
593
+ metrics = wrapped.get_savings_summary()
594
+ assert metrics["total_tokens_before"] >= metrics["total_tokens_after"]
595
+
596
+
597
+ # ============================================================================
598
+ # Test Class for optimize_messages standalone function
599
+ # ============================================================================
600
+
601
+
602
+ @pytest.mark.skipif(SKIP_BEDROCK, reason="AWS credentials not available")
603
+ @pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed")
604
+ class TestOptimizeMessagesFunction:
605
+ """Tests for the standalone optimize_messages function."""
606
+
607
+ def test_optimize_messages_basic(self):
608
+ """Test basic message optimization."""
609
+ from headroom.integrations.strands import optimize_messages
610
+
611
+ messages = [
612
+ {"role": "system", "content": "You are a helpful assistant."},
613
+ {"role": "user", "content": "Hello!"},
614
+ {"role": "assistant", "content": "Hi there! How can I help you today?"},
615
+ ]
616
+
617
+ optimized, metrics = optimize_messages(messages)
618
+
619
+ assert len(optimized) > 0
620
+
621
+ assert "tokens_before" in metrics
622
+ assert "tokens_after" in metrics
623
+ assert "tokens_saved" in metrics
624
+ assert metrics["tokens_before"] >= 0
625
+ assert metrics["tokens_after"] >= 0
626
+
627
+ def test_optimize_messages_with_tool_content(self):
628
+ """Test optimization of messages containing tool responses."""
629
+ from headroom.integrations.strands import optimize_messages
630
+
631
+ # Create messages with large tool output
632
+ large_data = json.dumps([{"id": i, "data": f"value_{i}" * 10} for i in range(100)])
633
+
634
+ messages = [
635
+ {"role": "system", "content": "You are a helpful assistant."},
636
+ {"role": "user", "content": "Get the data"},
637
+ {
638
+ "role": "assistant",
639
+ "content": None,
640
+ "tool_calls": [
641
+ {
642
+ "id": "call_123",
643
+ "type": "function",
644
+ "function": {"name": "get_data", "arguments": "{}"},
645
+ }
646
+ ],
647
+ },
648
+ {"role": "tool", "content": large_data, "tool_call_id": "call_123"},
649
+ {"role": "assistant", "content": "Here is the data summary..."},
650
+ ]
651
+
652
+ optimized, metrics = optimize_messages(messages)
653
+
654
+ assert len(optimized) > 0
655
+ assert metrics["tokens_before"] >= 0
656
+
657
+ def test_optimize_messages_custom_config(self):
658
+ """Test optimization with custom config."""
659
+ from headroom import HeadroomConfig
660
+ from headroom.integrations.strands import optimize_messages
661
+
662
+ config = HeadroomConfig()
663
+ config.smart_crusher.enabled = True
664
+ config.smart_crusher.min_tokens_to_crush = 10
665
+
666
+ messages = [
667
+ {"role": "user", "content": "Hello!"},
668
+ ]
669
+
670
+ optimized, metrics = optimize_messages(messages, config=config)
671
+
672
+ assert len(optimized) > 0
673
+ assert "tokens_before" in metrics
tests/integrations/test_strands/test_model_unit.py ADDED
@@ -0,0 +1,645 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Unit tests for Strands HeadroomStrandsModel.
2
+
3
+ These tests use mocks and do NOT require AWS credentials or strands-agents.
4
+ They test the internal logic of HeadroomStrandsModel in isolation.
5
+
6
+ For real integration tests, see test_model.py.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ from datetime import datetime, timezone
12
+ from unittest.mock import MagicMock, patch
13
+
14
+ import pytest
15
+
16
+ # Check if strands-agents is installed for proper skip handling
17
+ try:
18
+ import strands # noqa: F401
19
+
20
+ STRANDS_AVAILABLE = True
21
+ except ImportError:
22
+ STRANDS_AVAILABLE = False
23
+
24
+
25
+ # Skip all tests if Strands not installed
26
+ pytestmark = pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed")
27
+
28
+
29
+ # ============================================================================
30
+ # Fixtures
31
+ # ============================================================================
32
+
33
+
34
+ @pytest.fixture
35
+ def mock_strands_model():
36
+ """Create a mock Strands model."""
37
+ mock = MagicMock()
38
+ mock.config = {"model_id": "anthropic.claude-3-haiku-20240307-v1:0"}
39
+ mock.get_config.return_value = mock.config
40
+
41
+ # Mock the stream method as an async generator
42
+ async def mock_stream(*args, **kwargs):
43
+ yield {"type": "content", "data": "Hello"}
44
+ yield {"type": "content", "data": " world"}
45
+ yield {"type": "stop"}
46
+
47
+ mock.stream = mock_stream
48
+ return mock
49
+
50
+
51
+ @pytest.fixture
52
+ def sample_messages():
53
+ """Sample messages in Strands/OpenAI format."""
54
+ return [
55
+ {"role": "system", "content": "You are a helpful assistant."},
56
+ {"role": "user", "content": "What is the capital of France?"},
57
+ ]
58
+
59
+
60
+ @pytest.fixture
61
+ def large_conversation():
62
+ """Large conversation with many turns for compression testing."""
63
+ messages = [{"role": "system", "content": "You are a helpful assistant."}]
64
+ for i in range(50):
65
+ messages.append({"role": "user", "content": f"Question {i}: What is {i} + {i}?"})
66
+ messages.append({"role": "assistant", "content": f"The answer is {i + i}."})
67
+ return messages
68
+
69
+
70
+ # ============================================================================
71
+ # Test Classes
72
+ # ============================================================================
73
+
74
+
75
+ class TestHeadroomStrandsModelInit:
76
+ """Tests for HeadroomStrandsModel initialization."""
77
+
78
+ def test_init_with_defaults(self, mock_strands_model):
79
+ """Initialize with default settings."""
80
+ from headroom.integrations.strands import HeadroomStrandsModel
81
+
82
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
83
+
84
+ assert model.wrapped_model is mock_strands_model
85
+ assert model.total_tokens_saved == 0
86
+ assert model.metrics_history == []
87
+ assert model.auto_detect_provider is True
88
+
89
+ def test_init_with_custom_config(self, mock_strands_model):
90
+ """Initialize with custom HeadroomConfig."""
91
+ from headroom import HeadroomConfig
92
+ from headroom.integrations.strands import HeadroomStrandsModel
93
+
94
+ config = HeadroomConfig()
95
+ config.smart_crusher.min_tokens_to_crush = 100
96
+
97
+ model = HeadroomStrandsModel(
98
+ wrapped_model=mock_strands_model,
99
+ config=config,
100
+ auto_detect_provider=False,
101
+ )
102
+
103
+ assert model.headroom_config is config
104
+ assert model.auto_detect_provider is False
105
+
106
+ def test_init_requires_wrapped_model(self):
107
+ """Raises ValueError if wrapped_model is None."""
108
+ from headroom.integrations.strands import HeadroomStrandsModel
109
+
110
+ with pytest.raises(ValueError, match="wrapped_model cannot be None"):
111
+ HeadroomStrandsModel(wrapped_model=None)
112
+
113
+
114
+ class TestAttributeForwarding:
115
+ """Tests for attribute forwarding to wrapped model."""
116
+
117
+ def test_forwards_unknown_attributes(self, mock_strands_model):
118
+ """Forwards unknown attributes to wrapped model."""
119
+ from headroom.integrations.strands import HeadroomStrandsModel
120
+
121
+ mock_strands_model.custom_attr = "custom_value"
122
+ mock_strands_model.another_attr = 42
123
+
124
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
125
+
126
+ assert model.custom_attr == "custom_value"
127
+ assert model.another_attr == 42
128
+
129
+ def test_forwards_config_property(self, mock_strands_model):
130
+ """Forwards config property to wrapped model."""
131
+ from headroom.integrations.strands import HeadroomStrandsModel
132
+
133
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
134
+
135
+ config = model.config
136
+ assert config is mock_strands_model.config
137
+
138
+ def test_does_not_forward_internal_attrs(self, mock_strands_model):
139
+ """Does not forward internal wrapper attributes."""
140
+ from headroom.integrations.strands import HeadroomStrandsModel
141
+
142
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
143
+
144
+ # These should be wrapper's own attributes
145
+ assert model.wrapped_model is mock_strands_model
146
+ assert model.total_tokens_saved == 0
147
+ assert model.metrics_history == []
148
+
149
+ def test_get_config_delegates(self, mock_strands_model):
150
+ """get_config() delegates to wrapped model."""
151
+ from headroom.integrations.strands import HeadroomStrandsModel
152
+
153
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
154
+
155
+ config = model.get_config()
156
+ assert config == mock_strands_model.get_config()
157
+
158
+ def test_update_config_delegates(self, mock_strands_model):
159
+ """update_config() delegates to wrapped model."""
160
+ from headroom.integrations.strands import HeadroomStrandsModel
161
+
162
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
163
+
164
+ model.update_config(temperature=0.5)
165
+ mock_strands_model.update_config.assert_called_once_with(temperature=0.5)
166
+
167
+
168
+ class TestMessageConversion:
169
+ """Tests for message format conversion."""
170
+
171
+ def test_convert_dict_messages(self, mock_strands_model, sample_messages):
172
+ """Converts dict messages to OpenAI format."""
173
+ from headroom.integrations.strands import HeadroomStrandsModel
174
+
175
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
176
+
177
+ converted = model._convert_messages_to_openai(sample_messages)
178
+
179
+ assert len(converted) == 2
180
+ assert converted[0]["role"] == "system"
181
+ assert converted[0]["content"] == "You are a helpful assistant."
182
+ assert converted[1]["role"] == "user"
183
+ assert converted[1]["content"] == "What is the capital of France?"
184
+
185
+ def test_convert_messages_with_tool_calls(self, mock_strands_model):
186
+ """Converts messages with tool calls."""
187
+ from headroom.integrations.strands import HeadroomStrandsModel
188
+
189
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
190
+
191
+ messages = [
192
+ {
193
+ "role": "assistant",
194
+ "content": None,
195
+ "tool_calls": [
196
+ {"id": "call_123", "type": "function", "function": {"name": "search"}}
197
+ ],
198
+ },
199
+ {
200
+ "role": "tool",
201
+ "content": '{"results": []}',
202
+ "tool_call_id": "call_123",
203
+ "name": "search",
204
+ },
205
+ ]
206
+
207
+ converted = model._convert_messages_to_openai(messages)
208
+
209
+ assert len(converted) == 2
210
+ assert "tool_calls" in converted[0]
211
+ assert converted[1]["tool_call_id"] == "call_123"
212
+ assert converted[1]["name"] == "search"
213
+
214
+ def test_convert_message_objects(self, mock_strands_model):
215
+ """Converts message objects with role/content attributes."""
216
+ from headroom.integrations.strands import HeadroomStrandsModel
217
+
218
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
219
+
220
+ # Create mock message objects
221
+ msg1 = MagicMock()
222
+ msg1.role = "user"
223
+ msg1.content = "Hello"
224
+ msg1.tool_calls = None
225
+ msg1.tool_call_id = None
226
+ msg1.name = None
227
+
228
+ msg2 = MagicMock()
229
+ msg2.role = "assistant"
230
+ msg2.content = "Hi there!"
231
+ msg2.tool_calls = None
232
+ msg2.tool_call_id = None
233
+ msg2.name = None
234
+
235
+ converted = model._convert_messages_to_openai([msg1, msg2])
236
+
237
+ assert len(converted) == 2
238
+ assert converted[0]["role"] == "user"
239
+ assert converted[0]["content"] == "Hello"
240
+ assert converted[1]["role"] == "assistant"
241
+ assert converted[1]["content"] == "Hi there!"
242
+
243
+ def test_convert_handles_content_list(self, mock_strands_model):
244
+ """Converts messages with content as list (content blocks)."""
245
+ from headroom.integrations.strands import HeadroomStrandsModel
246
+
247
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
248
+
249
+ messages = [
250
+ {
251
+ "role": "user",
252
+ "content": [
253
+ {"type": "text", "text": "Look at this:"},
254
+ {"type": "image", "source": {"data": "base64..."}},
255
+ ],
256
+ }
257
+ ]
258
+
259
+ converted = model._convert_messages_to_openai(messages)
260
+
261
+ assert len(converted) == 1
262
+ assert isinstance(converted[0]["content"], list)
263
+ assert len(converted[0]["content"]) == 2
264
+
265
+
266
+ class TestOptimizeMessages:
267
+ """Tests for _optimize_messages method."""
268
+
269
+ def test_optimize_returns_metrics(self, mock_strands_model, sample_messages):
270
+ """_optimize_messages returns messages and metrics."""
271
+ from headroom.integrations.strands import HeadroomStrandsModel
272
+
273
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
274
+
275
+ # Mock the pipeline by setting _pipeline directly and mocking _headroom_provider
276
+ mock_pipeline = MagicMock()
277
+ mock_result = MagicMock()
278
+ mock_result.messages = sample_messages
279
+ mock_result.tokens_before = 50
280
+ mock_result.tokens_after = 40
281
+ mock_result.transforms_applied = ["cache_aligner"]
282
+ mock_pipeline.apply.return_value = mock_result
283
+
284
+ model._pipeline = mock_pipeline
285
+ model._headroom_provider = MagicMock()
286
+ model._headroom_provider.get_context_limit.return_value = 128000
287
+
288
+ optimized, metrics = model._optimize_messages(sample_messages)
289
+
290
+ assert len(optimized) == 2
291
+ assert metrics.tokens_before == 50
292
+ assert metrics.tokens_after == 40
293
+ assert metrics.tokens_saved == 10
294
+ assert "cache_aligner" in metrics.transforms_applied
295
+
296
+ def test_optimize_handles_empty_messages(self, mock_strands_model):
297
+ """_optimize_messages handles empty message list."""
298
+ from headroom.integrations.strands import HeadroomStrandsModel
299
+
300
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
301
+
302
+ optimized, metrics = model._optimize_messages([])
303
+
304
+ assert optimized == []
305
+ assert metrics.tokens_before == 0
306
+ assert metrics.tokens_after == 0
307
+ assert metrics.tokens_saved == 0
308
+
309
+ def test_optimize_tracks_metrics(self, mock_strands_model, sample_messages):
310
+ """_optimize_messages tracks metrics in history."""
311
+ from headroom.integrations.strands import HeadroomStrandsModel
312
+
313
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
314
+
315
+ # Mock the pipeline by setting _pipeline directly
316
+ mock_pipeline = MagicMock()
317
+ mock_result = MagicMock()
318
+ mock_result.messages = sample_messages
319
+ mock_result.tokens_before = 100
320
+ mock_result.tokens_after = 80
321
+ mock_result.transforms_applied = []
322
+ mock_pipeline.apply.return_value = mock_result
323
+
324
+ model._pipeline = mock_pipeline
325
+ model._headroom_provider = MagicMock()
326
+ model._headroom_provider.get_context_limit.return_value = 128000
327
+
328
+ model._optimize_messages(sample_messages)
329
+
330
+ assert len(model.metrics_history) == 1
331
+ assert model.metrics_history[0].tokens_saved == 20
332
+ assert model.total_tokens_saved == 20
333
+
334
+ def test_optimize_handles_pipeline_errors(self, mock_strands_model, sample_messages):
335
+ """_optimize_messages falls back on pipeline errors."""
336
+ from headroom.integrations.strands import HeadroomStrandsModel
337
+
338
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
339
+
340
+ # Mock the pipeline to raise an error
341
+ mock_pipeline = MagicMock()
342
+ mock_pipeline.apply.side_effect = ValueError("Pipeline error")
343
+
344
+ model._pipeline = mock_pipeline
345
+ model._headroom_provider = MagicMock()
346
+ model._headroom_provider.get_context_limit.return_value = 128000
347
+
348
+ # Should not raise, should fall back
349
+ optimized, metrics = model._optimize_messages(sample_messages)
350
+
351
+ assert len(optimized) == len(sample_messages)
352
+ assert "fallback:error" in metrics.transforms_applied
353
+
354
+
355
+ class TestPipelineLazyInit:
356
+ """Tests for TransformPipeline lazy initialization."""
357
+
358
+ def test_pipeline_is_lazily_initialized(self, mock_strands_model):
359
+ """Pipeline is not created until first access."""
360
+ from headroom.integrations.strands import HeadroomStrandsModel
361
+
362
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
363
+
364
+ # Should be None initially
365
+ assert model._pipeline is None
366
+
367
+ # Access pipeline property
368
+ with patch("headroom.integrations.strands.model.TransformPipeline"):
369
+ _ = model.pipeline
370
+
371
+ # Now should be initialized
372
+ assert model._pipeline is not None
373
+
374
+
375
+ class TestGetSavingsSummary:
376
+ """Tests for get_savings_summary method."""
377
+
378
+ def test_empty_summary(self, mock_strands_model):
379
+ """Returns zero values when no metrics recorded."""
380
+ from headroom.integrations.strands import HeadroomStrandsModel
381
+
382
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
383
+ summary = model.get_savings_summary()
384
+
385
+ assert summary["total_requests"] == 0
386
+ assert summary["total_tokens_saved"] == 0
387
+ assert summary["average_savings_percent"] == 0
388
+
389
+ def test_summary_with_metrics(self, mock_strands_model):
390
+ """Returns correct summary with recorded metrics."""
391
+ from headroom.integrations.strands import HeadroomStrandsModel
392
+ from headroom.integrations.strands.model import OptimizationMetrics
393
+
394
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
395
+
396
+ # Add metrics manually
397
+ model._metrics_history = [
398
+ OptimizationMetrics(
399
+ request_id="1",
400
+ timestamp=datetime.now(timezone.utc),
401
+ tokens_before=100,
402
+ tokens_after=80,
403
+ tokens_saved=20,
404
+ savings_percent=20.0,
405
+ transforms_applied=[],
406
+ model="test-model",
407
+ ),
408
+ OptimizationMetrics(
409
+ request_id="2",
410
+ timestamp=datetime.now(timezone.utc),
411
+ tokens_before=200,
412
+ tokens_after=120,
413
+ tokens_saved=80,
414
+ savings_percent=40.0,
415
+ transforms_applied=[],
416
+ model="test-model",
417
+ ),
418
+ ]
419
+ model._total_tokens_saved = 100
420
+
421
+ summary = model.get_savings_summary()
422
+
423
+ assert summary["total_requests"] == 2
424
+ assert summary["total_tokens_saved"] == 100
425
+ assert summary["average_savings_percent"] == 30.0 # (20 + 40) / 2
426
+ assert summary["total_tokens_before"] == 300
427
+ assert summary["total_tokens_after"] == 200
428
+
429
+
430
+ class TestReset:
431
+ """Tests for reset method."""
432
+
433
+ def test_reset_clears_all_state(self, mock_strands_model):
434
+ """reset() clears all tracked state."""
435
+ from headroom.integrations.strands import HeadroomStrandsModel
436
+ from headroom.integrations.strands.model import OptimizationMetrics
437
+
438
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
439
+
440
+ # Add some state
441
+ model._metrics_history = [
442
+ OptimizationMetrics(
443
+ request_id="1",
444
+ timestamp=datetime.now(timezone.utc),
445
+ tokens_before=100,
446
+ tokens_after=50,
447
+ tokens_saved=50,
448
+ savings_percent=50.0,
449
+ transforms_applied=[],
450
+ model="test",
451
+ )
452
+ ]
453
+ model._total_tokens_saved = 50
454
+
455
+ # Reset
456
+ model.reset()
457
+
458
+ # Verify all state cleared
459
+ assert model._metrics_history == []
460
+ assert model._total_tokens_saved == 0
461
+ assert model.total_tokens_saved == 0
462
+ assert len(model.metrics_history) == 0
463
+
464
+ # Summary should reflect reset
465
+ summary = model.get_savings_summary()
466
+ assert summary["total_requests"] == 0
467
+
468
+
469
+ class TestMetricsHistoryBound:
470
+ """Tests for metrics history bounding."""
471
+
472
+ def test_metrics_bounded_to_100(self, mock_strands_model):
473
+ """Metrics history is bounded to 100 entries."""
474
+ from headroom.integrations.strands import HeadroomStrandsModel
475
+ from headroom.integrations.strands.model import OptimizationMetrics
476
+
477
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
478
+
479
+ # Add 150 metrics
480
+ for i in range(150):
481
+ model._metrics_history.append(
482
+ OptimizationMetrics(
483
+ request_id=f"req_{i}",
484
+ timestamp=datetime.now(timezone.utc),
485
+ tokens_before=100,
486
+ tokens_after=80,
487
+ tokens_saved=20,
488
+ savings_percent=20.0,
489
+ transforms_applied=[],
490
+ model="test",
491
+ )
492
+ )
493
+ # Simulate what _optimize_messages does
494
+ if len(model._metrics_history) > 100:
495
+ model._metrics_history = model._metrics_history[-100:]
496
+
497
+ # Should be bounded at 100
498
+ assert len(model.metrics_history) == 100
499
+
500
+ # Should contain the most recent entries
501
+ assert model.metrics_history[-1].request_id == "req_149"
502
+
503
+
504
+ class TestOptimizeMessagesFunction:
505
+ """Tests for standalone optimize_messages function."""
506
+
507
+ def test_optimize_messages_basic(self):
508
+ """optimize_messages processes messages and returns metrics."""
509
+ from headroom.integrations.strands import optimize_messages
510
+
511
+ messages = [
512
+ {"role": "user", "content": "Hello"},
513
+ {"role": "assistant", "content": "Hi there!"},
514
+ ]
515
+
516
+ with patch("headroom.integrations.strands.model.TransformPipeline") as MockPipeline:
517
+ mock_instance = MagicMock()
518
+ mock_result = MagicMock()
519
+ mock_result.messages = messages
520
+ mock_result.tokens_before = 20
521
+ mock_result.tokens_after = 15
522
+ mock_result.transforms_applied = ["cache_aligner"]
523
+ mock_instance.apply.return_value = mock_result
524
+ MockPipeline.return_value = mock_instance
525
+
526
+ optimized, metrics = optimize_messages(messages)
527
+
528
+ assert len(optimized) == 2
529
+ assert metrics["tokens_saved"] == 5
530
+ assert metrics["savings_percent"] == 25.0
531
+
532
+ def test_optimize_messages_with_custom_config(self):
533
+ """optimize_messages uses custom config."""
534
+ from headroom import HeadroomConfig
535
+ from headroom.integrations.strands import optimize_messages
536
+
537
+ config = HeadroomConfig()
538
+ messages = [{"role": "user", "content": "Test"}]
539
+
540
+ with patch("headroom.integrations.strands.model.TransformPipeline") as MockPipeline:
541
+ mock_instance = MagicMock()
542
+ mock_result = MagicMock()
543
+ mock_result.messages = messages
544
+ mock_result.tokens_before = 10
545
+ mock_result.tokens_after = 10
546
+ mock_result.transforms_applied = []
547
+ mock_instance.apply.return_value = mock_result
548
+ MockPipeline.return_value = mock_instance
549
+
550
+ optimized, metrics = optimize_messages(messages, config=config)
551
+
552
+ # Verify config was passed to pipeline
553
+ MockPipeline.assert_called_once()
554
+ call_kwargs = MockPipeline.call_args[1]
555
+ assert call_kwargs["config"] is config
556
+
557
+
558
+ class TestStreamMethod:
559
+ """Tests for stream method."""
560
+
561
+ @pytest.mark.asyncio
562
+ async def test_stream_optimizes_messages(self, mock_strands_model, sample_messages):
563
+ """stream() applies optimization before calling wrapped model."""
564
+ from headroom.integrations.strands import HeadroomStrandsModel
565
+
566
+ model = HeadroomStrandsModel(wrapped_model=mock_strands_model)
567
+
568
+ # Mock the optimization
569
+ with patch.object(model, "_optimize_messages") as mock_optimize:
570
+ mock_optimize.return_value = (
571
+ sample_messages,
572
+ MagicMock(
573
+ tokens_before=50,
574
+ tokens_after=40,
575
+ savings_percent=20.0,
576
+ ),
577
+ )
578
+
579
+ # Consume the stream
580
+ events = []
581
+ async for event in model.stream(sample_messages):
582
+ events.append(event)
583
+
584
+ # Should have called optimization
585
+ mock_optimize.assert_called_once()
586
+
587
+ # Should have yielded events from wrapped model
588
+ assert len(events) > 0
589
+
590
+
591
+ class TestStrandsAvailableFunction:
592
+ """Tests for strands_available function."""
593
+
594
+ def test_strands_available_returns_bool(self):
595
+ """strands_available() returns boolean."""
596
+ from headroom.integrations.strands import strands_available
597
+
598
+ result = strands_available()
599
+
600
+ # Since we're in a test where strands is available (skipif passed)
601
+ assert isinstance(result, bool)
602
+ assert result is True
603
+
604
+
605
+ class TestRealHeadroomIntegration:
606
+ """Integration tests with real Headroom (no mocking)."""
607
+
608
+ def test_real_optimization_with_mock_model(self, mock_strands_model, sample_messages):
609
+ """Test with real Headroom transforms (no API calls)."""
610
+ from headroom.integrations.strands import HeadroomStrandsModel
611
+
612
+ model = HeadroomStrandsModel(
613
+ wrapped_model=mock_strands_model,
614
+ auto_detect_provider=False, # Use default OpenAI provider
615
+ )
616
+
617
+ # This calls real Headroom optimization
618
+ optimized, metrics = model._optimize_messages(sample_messages)
619
+
620
+ # Should return valid messages
621
+ assert len(optimized) >= 1
622
+ assert all("role" in m and "content" in m for m in optimized)
623
+
624
+ # Metrics should be tracked
625
+ assert len(model.metrics_history) == 1
626
+ assert metrics.tokens_before >= 0
627
+ assert metrics.tokens_after >= 0
628
+
629
+ def test_large_conversation_handling(self, mock_strands_model, large_conversation):
630
+ """Large conversations are processed without errors."""
631
+ from headroom.integrations.strands import HeadroomStrandsModel
632
+
633
+ model = HeadroomStrandsModel(
634
+ wrapped_model=mock_strands_model,
635
+ auto_detect_provider=False,
636
+ )
637
+
638
+ # Should handle large conversation without errors
639
+ optimized, metrics = model._optimize_messages(large_conversation)
640
+
641
+ # Should return messages
642
+ assert len(optimized) >= 1
643
+
644
+ # Metrics should show processing occurred
645
+ assert metrics.tokens_before > 0