Spaces:
Build error
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 +56 -0
- examples/strands_bedrock_demo.py +1001 -0
- headroom/integrations/strands/__init__.py +88 -0
- headroom/integrations/strands/hooks.py +540 -0
- headroom/integrations/strands/model.py +625 -0
- headroom/integrations/strands/providers.py +166 -0
- pyproject.toml +4 -0
- tests/integrations/test_strands/__init__.py +1 -0
- tests/integrations/test_strands/test_hooks.py +545 -0
- tests/integrations/test_strands/test_hooks_unit.py +564 -0
- tests/integrations/test_strands/test_model.py +673 -0
- tests/integrations/test_strands/test_model_unit.py +645 -0
|
@@ -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.
|
|
@@ -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())
|
|
@@ -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 |
+
]
|
|
@@ -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")
|
|
@@ -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
|
|
@@ -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"
|
|
@@ -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
|
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
"""Tests for Strands Agents SDK integration with Headroom."""
|
|
@@ -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
|
|
@@ -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"}]
|
|
@@ -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
|
|
@@ -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
|