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

feat: add CrewAI and AutoGen tool compression integrations (#1384)

Browse source 
## Description

Add CrewAI and AutoGen tool compression integrations, following the same
patterns as the existing LangChain agent integration
(`HeadroomToolWrapper` / `wrap_tools_with_headroom`). Both delegate
compression to `compress_tool_result()` from the MCP integration, with
per-tool metrics tracking via `ToolCompressionMetrics` /
`ToolMetricsCollector`.

Closes #1379

## Type of Change

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

## Changes Made

- Add `headroom/integrations/crewai/` — `HeadroomToolWrapper` subclasses
CrewAI `BaseTool`, wraps `_run()` with compression
- Add `headroom/integrations/autogen/` — `HeadroomToolWrapper` wraps
AutoGen `FunctionTool` (sync and async) with compression
- Wire both into `headroom/integrations/__init__.py` with aliased
re-exports (avoids name collision with LangChain's
`HeadroomToolWrapper`)
- Add `[crewai]` and `[autogen]` optional dependency extras to
`pyproject.toml`
- Add 24 unit tests (12 per framework) under `tests/test_integrations/`
- Add `.mdx` doc pages for both frameworks under `docs/content/docs/`
- Update `CHANGELOG.md` with entries under `### Added`

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [ ] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ ruff check headroom/integrations/crewai headroom/integrations/autogen tests/test_integrations/crewai tests/test_integrations/autogen
All checks passed!

$ pytest tests/test_integrations/autogen -v
12 passed

$ pytest tests/test_integrations/crewai -v
12 passed
```

## Real Behavior Proof

- Environment: Windows 11, Python 3.11, crewai 1.14.7, autogen-agentchat
0.7.5
- Exact command / steps: Ran standalone adapter demos and benchmark
runner across 4 task types
- Observed result:

| Task | Tokens (raw) | Tokens (compressed) | Savings |
|------|-------------|-------------------|---------|
| Inventory JSON (80 items) | 5,044 | 1,532 | 69.6% |
| Server logs (150 lines) | 8,712 | 314 | 96.4% |
| Analytics query (100 rows) | 10,762 | 10,762 | 0% |
| API docs (20 endpoints) | 8,043 | 8,043 | 0% |

Compression results are identical across CrewAI and AutoGen — expected
since both route through the same `compress_tool_result()` pipeline.

- Not tested: Full end-to-end with a live LLM agent loop (demos test the
compression pipeline standalone). LangGraph not included — headroom
already has `headroom/integrations/langchain/langgraph.py`.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable

## Additional Notes

- LangGraph integration is intentionally excluded — headroom already has
one at `headroom/integrations/langchain/langgraph.py`
- Re-exports in `__init__.py` are aliased (`CrewAIToolWrapper`,
`AutoGenToolWrapper`) to avoid collision with the existing LangChain
`HeadroomToolWrapper`
- Both integrations follow the exact same conventions as the existing
LangChain agents module: optional dep guard, `compress_tool_result()`
delegation, metrics with 1000-entry cap, Google-style docstrings
- `mypy` not checked due to Rust build dependency (`maturin`) that
requires Application Control policy changes on this machine

---------

Co-authored-by: Sneha27feb <sroy27.ai@gmail.com>
This commit is contained in:
Sneha Roy 2026-07-16 01:28:54 +05:30 • committed by GitHub
parent 3dd9660d91
commit e8bff1cfe3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 1698 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

5
CHANGELOG.md
View file

@ -102,6 +102,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
* **docs/ccr:** qualify the current CCR auto-resolution claim by provider. The docs now state that transparent `headroom_retrieve` handling is wired on the Anthropic and OpenAI proxy paths, while native Gemini still lacks that server-side response-handler path and Gemini's OpenAI-compatible endpoint can fail round-2 continuations with `MALFORMED_FUNCTION_CALL` ([#2041](https://github.com/headroomlabs-ai/headroom/issues/2041)).
* **docs/claude:** document that `ENABLE_TOOL_SEARCH=true` is correct for the standalone Claude CLI through Headroom but currently breaks tool-result rendering in Anthropic's VSCode extension webview, and point persistent-install users at the manifest override to set `tool_envs.claude.ENABLE_TOOL_SEARCH` to `"false"` for that target ([#2028](https://github.com/headroomlabs-ai/headroom/issues/2028)).
### Added
* **integrations:** CrewAI tool compression — `wrap_tools_with_headroom()` wraps CrewAI `BaseTool` instances with automatic output compression via `compress_tool_result()`, with per-tool metrics tracking ([#1379](https://github.com/headroomlabs-ai/headroom/issues/1379)).
* **integrations:** AutoGen tool compression — `wrap_tools_with_headroom()` wraps AutoGen `FunctionTool` instances (sync and async) with automatic output compression, including per-tool metrics tracking ([#1379](https://github.com/headroomlabs-ai/headroom/issues/1379)).
### Features
* **wrap:** add `headroom wrap omp` / `headroom unwrap omp` for Oh My Pi — points omp's built-in `anthropic` provider at the local proxy via a marker-fenced `providers.anthropic.baseUrl` override in `~/.omp/agent/models.yml`, snapshotting the pre-wrap file byte-for-byte and restoring it on unwrap. omp resolves its Anthropic chat endpoint from models.yml (`ANTHROPIC_BASE_URL` only feeds its web-search helper), and a same-ID override keeps omp's bundled model catalog and stored credentials ([#1149](https://github.com/headroomlabs-ai/headroom/issues/1149))

150
docs/content/docs/autogen.mdx Normal file
View file

@ -0,0 +1,150 @@
---
title: AutoGen
description: Automatic tool output compression for AutoGen agents with per-tool metrics tracking.
---
Headroom integrates with [AutoGen](https://github.com/microsoft/autogen) (`autogen-agentchat` >=0.7) to compress tool outputs before they enter the agent's model context. Tool-heavy agents that return large JSON arrays, database results, or verbose logs see 60-90% token reduction.
## Installation
```bash
pip install headroom-ai autogen-agentchat
```
## Quick start
Wrap tools in one line:
```python
from autogen_agentchat.agents import AssistantAgent
from autogen_core.tools import FunctionTool
from headroom.integrations.autogen import wrap_tools_with_headroom
def search_database(query: str) -> str:
"""Search the database and return results."""
return json.dumps({"results": [...], "total": 1000})
tool = FunctionTool(search_database, description="Search the database")
wrapped = wrap_tools_with_headroom([tool])
agent = AssistantAgent(
name="researcher",
model_client=model_client,
tools=wrapped,
)
```
## Per-tool metrics
Track compression stats across all tool invocations:
```python
from headroom.integrations.autogen import get_tool_metrics
metrics = get_tool_metrics()
print(metrics.get_summary())
# {
# 'total_invocations': 25,
# 'total_compressions': 18,
# 'total_chars_saved': 450000,
# 'average_compression_ratio': 0.35,
# 'by_tool': {
# 'search_database': {'invocations': 15, 'compressions': 12, 'chars_saved': 320000},
# }
# }
```
Reset between sessions:
```python
from headroom.integrations.autogen import reset_tool_metrics
reset_tool_metrics()
```
## Custom configuration
Control the compression threshold:
```python
wrapped = wrap_tools_with_headroom(
[search_tool, log_tool],
min_chars_to_compress=500, # Default: 1000
)
```
Use a dedicated metrics collector:
```python
from headroom.integrations.autogen import ToolMetricsCollector, wrap_tools_with_headroom
collector = ToolMetricsCollector()
wrapped = wrap_tools_with_headroom(
[search_tool],
metrics_collector=collector,
)
print(collector.get_summary())
```
## Wrapping individual tools
For finer control, wrap tools individually:
```python
from headroom.integrations.autogen import HeadroomToolWrapper
wrapper = HeadroomToolWrapper(
search_tool,
min_chars_to_compress=500,
)
# Get the wrapped FunctionTool
compressed_tool = wrapper.as_function_tool()
agent = AssistantAgent(
name="researcher",
model_client=model_client,
tools=[compressed_tool],
)
```
## Async support
AutoGen tools are natively async. The wrapper handles both sync and async
tool functions transparently:
```python
async def async_search(query: str) -> str:
"""Async database search."""
results = await db.search(query)
return json.dumps(results)
tool = FunctionTool(async_search, description="Async search")
wrapped = wrap_tools_with_headroom([tool])
# Compression works identically for async tools
```
## How it works
AutoGen routes tool execution through `FunctionTool`, which wraps a plain
Python function. The function's return value is stringified and becomes
`FunctionExecutionResult.content` — what the LLM reads on its next turn.
`HeadroomToolWrapper` creates a new `FunctionTool` with a wrapper function that:
1. Calls the original function
2. Checks if the stringified output exceeds `min_chars_to_compress`
3. If so, compresses via Headroom's `compress_tool_result()`
4. Records metrics and returns the compressed string
The wrapper preserves the original tool's name, description, and parameter
schema, so it works as a drop-in replacement.
### Why not `tool_call_summary_formatter`?
AutoGen's `AssistantAgent` accepts a `tool_call_summary_formatter` parameter,
which looks like a natural hook. However, it only controls the **final summary
message** emitted after the tool loop exits — it does not touch the raw
`FunctionExecutionResult` that gets added to `model_context` (what the LLM
actually reads). Wrapping the function is the only clean interception point.

123
docs/content/docs/crewai.mdx Normal file
View file

@ -0,0 +1,123 @@
---
title: CrewAI
description: Automatic tool output compression for CrewAI agents with per-tool metrics tracking.
---
Headroom integrates with [CrewAI](https://github.com/crewAIInc/crewAI) to compress tool outputs before they enter the agent's LLM context. Tool-heavy agents that return large JSON arrays, database results, or verbose logs see 60-90% token reduction.
## Installation
```bash
pip install headroom-ai crewai
```
## Quick start
Wrap tools in one line:
```python
from crewai import Agent, Crew, Task
from crewai.tools.base_tool import tool
from headroom.integrations.crewai import wrap_tools_with_headroom
@tool
def search_database(query: str) -> str:
"""Search the database and return results."""
return json.dumps({"results": [...], "total": 1000})
wrapped = wrap_tools_with_headroom([search_database])
agent = Agent(
role="Researcher",
goal="Answer questions using data",
backstory="You research things.",
tools=wrapped,
)
task = Task(description="Find all active users", agent=agent, expected_output="Summary")
crew = Crew(agents=[agent], tasks=[task])
crew.kickoff()
```
## Per-tool metrics
Track compression stats across all tool invocations:
```python
from headroom.integrations.crewai import get_tool_metrics
metrics = get_tool_metrics()
print(metrics.get_summary())
# {
# 'total_invocations': 25,
# 'total_compressions': 18,
# 'total_chars_saved': 450000,
# 'average_compression_ratio': 0.35,
# 'by_tool': {
# 'search_database': {'invocations': 15, 'compressions': 12, 'chars_saved': 320000},
# 'fetch_logs': {'invocations': 10, 'compressions': 6, 'chars_saved': 130000},
# }
# }
```
Reset between sessions:
```python
from headroom.integrations.crewai import reset_tool_metrics
reset_tool_metrics()
```
## Custom configuration
Control the compression threshold:
```python
wrapped = wrap_tools_with_headroom(
[search_database, fetch_logs],
min_chars_to_compress=500, # Default: 1000
)
```
Use a dedicated metrics collector instead of the global one:
```python
from headroom.integrations.crewai import ToolMetricsCollector, wrap_tools_with_headroom
collector = ToolMetricsCollector()
wrapped = wrap_tools_with_headroom(
[search_database],
metrics_collector=collector,
)
# After crew run
print(collector.get_summary())
```
## Wrapping individual tools
For finer control, wrap tools individually:
```python
from headroom.integrations.crewai import HeadroomToolWrapper
wrapper = HeadroomToolWrapper(
search_database,
min_chars_to_compress=500,
)
# Use wrapper directly — it's a BaseTool
agent = Agent(role="Researcher", tools=[wrapper], ...)
```
## How it works
CrewAI tools extend `BaseTool` with a `run()` → `_run()` execution flow.
`HeadroomToolWrapper` subclasses `BaseTool` and overrides `_run()` to:
1. Call the original tool's `run()` method
2. Check if the output exceeds `min_chars_to_compress`
3. If so, compress via Headroom's `compress_tool_result()`
4. Record metrics and return the compressed output
The wrapper preserves the original tool's name, description, and argument
schema, so it works as a drop-in replacement anywhere CrewAI expects a tool.

72
headroom/integrations/__init__.py
View file

@ -15,6 +15,14 @@ Agno (pip install agno):
- HeadroomPreHook/HeadroomPostHook: Agent-level hooks for tracking
- create_headroom_hooks: Convenience function to create hook pairs
CrewAI (pip install headroom[crewai]):
- HeadroomToolWrapper: Tool output compression for CrewAI agents
- wrap_tools_with_headroom: Batch wrapper for CrewAI tools
AutoGen (pip install headroom[autogen]):
- HeadroomToolWrapper: Tool output compression for AutoGen agents
- wrap_tools_with_headroom: Batch wrapper for AutoGen tools
MCP (Model Context Protocol):
- HeadroomMCPCompressor: Compress MCP tool results
- compress_tool_result: Simple function for tool compression
@ -103,6 +111,56 @@ try:
except ImportError:
_AGNO_AVAILABLE = False
# Re-export from crewai subpackage (optional dependency)
try:
from .crewai import (
HeadroomToolWrapper as CrewAIToolWrapper,
)
from .crewai import (
ToolCompressionMetrics as CrewAIToolCompressionMetrics,
)
from .crewai import (
ToolMetricsCollector as CrewAIToolMetricsCollector,
)
from .crewai import (
get_tool_metrics as get_crewai_tool_metrics,
)
from .crewai import (
reset_tool_metrics as reset_crewai_tool_metrics,
)
from .crewai import (
wrap_tools_with_headroom as wrap_crewai_tools,
)
_CREWAI_AVAILABLE = True
except ImportError:
_CREWAI_AVAILABLE = False
# Re-export from autogen subpackage (optional dependency)
try:
from .autogen import (
HeadroomToolWrapper as AutoGenToolWrapper,
)
from .autogen import (
ToolCompressionMetrics as AutoGenToolCompressionMetrics,
)
from .autogen import (
ToolMetricsCollector as AutoGenToolMetricsCollector,
)
from .autogen import (
get_tool_metrics as get_autogen_tool_metrics,
)
from .autogen import (
reset_tool_metrics as reset_autogen_tool_metrics,
)
from .autogen import (
wrap_tools_with_headroom as wrap_autogen_tools,
)
_AUTOGEN_AVAILABLE = True
except ImportError:
_AUTOGEN_AVAILABLE = False
__all__ = [
# LangChain Core
"HeadroomChatModel",
@ -156,4 +214,18 @@ __all__ = [
"get_model_name_from_agno",
"AgnoOptimizationMetrics",
"optimize_agno_messages",
# CrewAI
"CrewAIToolWrapper",
"CrewAIToolCompressionMetrics",
"CrewAIToolMetricsCollector",
"wrap_crewai_tools",
"get_crewai_tool_metrics",
"reset_crewai_tool_metrics",
# AutoGen
"AutoGenToolWrapper",
"AutoGenToolCompressionMetrics",
"AutoGenToolMetricsCollector",
"wrap_autogen_tools",
"get_autogen_tool_metrics",
"reset_autogen_tool_metrics",
]

45
headroom/integrations/autogen/__init__.py Normal file
View file

@ -0,0 +1,45 @@
"""AutoGen integration for Headroom.
This module provides tool output compression for AutoGen agents,
wrapping FunctionTool instances so their outputs are automatically
compressed before entering the agent's model context.
Components:
- HeadroomToolWrapper: Wraps a single AutoGen FunctionTool with compression
- wrap_tools_with_headroom: Wraps multiple tools at once
- ToolCompressionMetrics: Per-invocation metrics dataclass
- ToolMetricsCollector: Aggregates metrics across all invocations
Example:
from autogen_agentchat.agents import AssistantAgent
from autogen_core.tools import FunctionTool
from headroom.integrations.autogen import wrap_tools_with_headroom
def search_db(query: str) -> str:
return json.dumps(results)
tool = FunctionTool(search_db, description="Search the database")
wrapped = wrap_tools_with_headroom([tool])
agent = AssistantAgent(name="researcher", tools=wrapped, ...)
Install: pip install headroom-ai autogen-agentchat
"""
from .agents import (
HeadroomToolWrapper,
ToolCompressionMetrics,
ToolMetricsCollector,
get_tool_metrics,
reset_tool_metrics,
wrap_tools_with_headroom,
)
__all__ = [
"HeadroomToolWrapper",
"ToolCompressionMetrics",
"ToolMetricsCollector",
"wrap_tools_with_headroom",
"get_tool_metrics",
"reset_tool_metrics",
]

386
headroom/integrations/autogen/agents.py Normal file
View file

@ -0,0 +1,386 @@
"""AutoGen agent tool integration with output compression.
This module provides HeadroomToolWrapper and wrap_tools_with_headroom
for wrapping AutoGen FunctionTool instances to automatically compress
their outputs and track per-tool compression metrics.
AutoGen (autogen-agentchat >=0.7) routes tool execution through a
Workbench abstraction. FunctionTool wraps a plain Python function;
the function's return value is stringified and becomes the
FunctionExecutionResult.content that enters model_context.
Interception strategy: wrap the callable inside FunctionTool so the
return value is compressed before AutoGen stringifies it. This is
the same pattern as the LangChain/CrewAI tool wrappers.
Note: AutoGen's ``tool_call_summary_formatter`` parameter on
AssistantAgent only controls the *final summary* emitted after the
tool loop, not what enters model_context. Wrapping the function
is the only clean, version-stable hook.
Example:
from autogen_core.tools import FunctionTool
from headroom.integrations.autogen import wrap_tools_with_headroom
def search_database(query: str) -> str:
\"\"\"Search the database.\"\"\"
return json.dumps({"results": [...], "total": 1000})
tool = FunctionTool(search_database, description="Search")
wrapped = wrap_tools_with_headroom([tool])
"""
from __future__ import annotations
import asyncio
import functools
import logging
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any
try:
from autogen_core.tools import FunctionTool
AUTOGEN_AVAILABLE = True
except ImportError:
AUTOGEN_AVAILABLE = False
FunctionTool = object # type: ignore[misc,assignment]
from headroom.integrations.mcp import compress_tool_result
logger = logging.getLogger(__name__)
def _check_autogen_available() -> None:
"""Raise ImportError if AutoGen is not installed."""
if not AUTOGEN_AVAILABLE:
raise ImportError(
"AutoGen is required for this integration. Install with: pip install autogen-agentchat"
)
@dataclass
class ToolCompressionMetrics:
"""Metrics from a single tool compression.
Attributes:
tool_name: Name of the tool that was invoked.
timestamp: When the compression occurred.
chars_before: Character count of the original output.
chars_after: Character count after compression.
chars_saved: Characters removed by compression.
compression_ratio: Ratio of compressed to original size.
was_compressed: Whether compression was actually applied.
"""
tool_name: str
timestamp: datetime
chars_before: int
chars_after: int
chars_saved: int
compression_ratio: float
was_compressed: bool
@dataclass
class ToolMetricsCollector:
"""Collects compression metrics across all tool invocations.
Attributes:
metrics: List of per-invocation metrics.
"""
metrics: list[ToolCompressionMetrics] = field(default_factory=list)
def add(self, metric: ToolCompressionMetrics) -> None:
"""Add a metric entry.
Args:
metric: The compression metrics to record.
"""
self.metrics.append(metric)
if len(self.metrics) > 1000:
self.metrics = self.metrics[-1000:]
def get_summary(self) -> dict[str, Any]:
"""Get summary statistics.
Returns:
Dict with total_invocations, total_compressions,
total_chars_saved, average_compression_ratio, and
per-tool breakdown.
"""
if not self.metrics:
return {
"total_invocations": 0,
"total_compressions": 0,
"total_chars_saved": 0,
}
compressed = [m for m in self.metrics if m.was_compressed]
return {
"total_invocations": len(self.metrics),
"total_compressions": len(compressed),
"total_chars_saved": sum(m.chars_saved for m in self.metrics),
"average_compression_ratio": (
sum(m.compression_ratio for m in compressed) / len(compressed) if compressed else 0
),
"by_tool": self._get_by_tool_stats(),
}
def _get_by_tool_stats(self) -> dict[str, dict[str, Any]]:
"""Get per-tool statistics."""
by_tool: dict[str, list[ToolCompressionMetrics]] = {}
for m in self.metrics:
if m.tool_name not in by_tool:
by_tool[m.tool_name] = []
by_tool[m.tool_name].append(m)
result = {}
for name, tool_metrics in by_tool.items():
compressed = [m for m in tool_metrics if m.was_compressed]
result[name] = {
"invocations": len(tool_metrics),
"compressions": len(compressed),
"chars_saved": sum(m.chars_saved for m in tool_metrics),
}
return result
# Global metrics collector
_global_metrics = ToolMetricsCollector()
def get_tool_metrics() -> ToolMetricsCollector:
"""Get the global tool metrics collector.
Returns:
The global ToolMetricsCollector instance.
"""
return _global_metrics
def reset_tool_metrics() -> None:
"""Reset global tool metrics."""
global _global_metrics
_global_metrics = ToolMetricsCollector()
def _compress_and_record(
output: str,
tool_name: str,
min_chars: int,
metrics: ToolMetricsCollector,
) -> str:
"""Compress output and record metrics.
Args:
output: Tool output string.
tool_name: Name of the tool for logging.
min_chars: Minimum chars to trigger compression.
metrics: Collector for metrics.
Returns:
Compressed output, or original if below threshold or on error.
"""
chars_before = len(output)
if chars_before < min_chars:
_record_metrics(metrics, tool_name, output, output, was_compressed=False)
return output
try:
compressed = compress_tool_result(
content=output,
tool_name=tool_name,
)
except Exception as e:
logger.debug("Tool compression failed for %s: %s", tool_name, e)
_record_metrics(metrics, tool_name, output, output, was_compressed=False)
return output
_record_metrics(metrics, tool_name, output, compressed, was_compressed=True)
return compressed
def _record_metrics(
collector: ToolMetricsCollector,
tool_name: str,
original: str,
compressed: str,
was_compressed: bool,
) -> None:
"""Record compression metrics.
Args:
collector: The metrics collector.
tool_name: Name of the tool.
original: Original output.
compressed: Compressed output.
was_compressed: Whether compression was applied.
"""
chars_before = len(original)
chars_after = len(compressed)
chars_saved = chars_before - chars_after
metric = ToolCompressionMetrics(
tool_name=tool_name,
timestamp=datetime.now(),
chars_before=chars_before,
chars_after=chars_after,
chars_saved=max(0, chars_saved),
compression_ratio=chars_after / chars_before if chars_before > 0 else 1.0,
was_compressed=was_compressed and chars_saved > 0,
)
collector.add(metric)
if was_compressed and chars_saved > 0:
logger.info(
"HeadroomToolWrapper[%s]: %d -> %d chars (%d saved, %.1f%% of original)",
tool_name,
chars_before,
chars_after,
chars_saved,
metric.compression_ratio * 100,
)
class HeadroomToolWrapper:
"""Wraps an AutoGen FunctionTool to compress its output.
Creates a new FunctionTool whose internal function calls the original,
stringifies the result, compresses it, and returns the compressed string.
The original tool's name, description, and parameter schema are preserved.
Example:
from autogen_core.tools import FunctionTool
from headroom.integrations.autogen import HeadroomToolWrapper
def search(query: str) -> str:
return json.dumps({"results": [...]})
tool = FunctionTool(search, description="Search")
wrapper = HeadroomToolWrapper(tool)
wrapped_tool = wrapper.as_function_tool()
Attributes:
name: Tool name (from wrapped tool).
description: Tool description (from wrapped tool).
wrapped_tool: The new FunctionTool with compression.
"""
def __init__(
self,
tool: FunctionTool,
min_chars_to_compress: int = 1000,
metrics_collector: ToolMetricsCollector | None = None,
) -> None:
"""Initialize HeadroomToolWrapper.
Args:
tool: The AutoGen FunctionTool to wrap.
min_chars_to_compress: Minimum character count for output
before compression is applied. Default 1000.
metrics_collector: Collector for metrics. Uses global
collector if not specified.
"""
_check_autogen_available()
self.name = tool.name
self.description = tool.description
self._min_chars = min_chars_to_compress
self._metrics = metrics_collector or _global_metrics
self.wrapped_tool = self._create_wrapped_tool(tool)
def _create_wrapped_tool(self, tool: FunctionTool) -> FunctionTool:
"""Create a new FunctionTool with compression.
Args:
tool: The original FunctionTool.
Returns:
A new FunctionTool that compresses output.
"""
original_func = tool._func
tool_name = tool.name
min_chars = self._min_chars
metrics = self._metrics
if asyncio.iscoroutinefunction(original_func):
@functools.wraps(original_func)
async def _compressed_func(*args: Any, **kwargs: Any) -> str:
raw = await original_func(*args, **kwargs)
return _compress_and_record(str(raw), tool_name, min_chars, metrics)
else:
@functools.wraps(original_func)
def _compressed_func(*args: Any, **kwargs: Any) -> str:
raw = original_func(*args, **kwargs)
return _compress_and_record(str(raw), tool_name, min_chars, metrics)
return FunctionTool(
_compressed_func,
description=tool.description,
name=tool.name,
)
def as_function_tool(self) -> FunctionTool:
"""Return the wrapped FunctionTool.
Returns:
FunctionTool with compression applied.
"""
return self.wrapped_tool
def wrap_tools_with_headroom(
tools: list[FunctionTool],
min_chars_to_compress: int = 1000,
metrics_collector: ToolMetricsCollector | None = None,
) -> list[FunctionTool]:
"""Wrap multiple AutoGen FunctionTools with Headroom compression.
Convenience function to wrap all tools in a list at once.
Each wrapped tool preserves the original's name, description,
and parameter schema.
Args:
tools: List of AutoGen FunctionTools to wrap.
min_chars_to_compress: Minimum output size for compression.
metrics_collector: Shared metrics collector for all tools.
Returns:
List of wrapped FunctionTools.
Example:
from autogen_agentchat.agents import AssistantAgent
from autogen_core.tools import FunctionTool
from headroom.integrations.autogen import wrap_tools_with_headroom
def search(query: str) -> str:
return json.dumps(results)
tool = FunctionTool(search, description="Search")
wrapped = wrap_tools_with_headroom([tool])
agent = AssistantAgent(
name="researcher",
model_client=model_client,
tools=wrapped,
)
"""
_check_autogen_available()
collector = metrics_collector or _global_metrics
return [
HeadroomToolWrapper(
tool=t,
min_chars_to_compress=min_chars_to_compress,
metrics_collector=collector,
).as_function_tool()
for t in tools
]

45
headroom/integrations/crewai/__init__.py Normal file
View file

@ -0,0 +1,45 @@
"""CrewAI integration for Headroom.
This module provides tool output compression for CrewAI agents,
wrapping BaseTool instances so their outputs are automatically
compressed before entering the agent's LLM context.
Components:
- HeadroomToolWrapper: Wraps a single CrewAI BaseTool with compression
- wrap_tools_with_headroom: Wraps multiple tools at once
- ToolCompressionMetrics: Per-invocation metrics dataclass
- ToolMetricsCollector: Aggregates metrics across all invocations
Example:
from crewai import Agent, Crew, Task
from crewai.tools.base_tool import tool
from headroom.integrations.crewai import wrap_tools_with_headroom
@tool
def search_db(query: str) -> str:
\"\"\"Search the database.\"\"\"
return json.dumps(results)
wrapped = wrap_tools_with_headroom([search_db])
agent = Agent(role="Researcher", tools=wrapped, ...)
Install: pip install headroom-ai crewai
"""
from .agents import (
HeadroomToolWrapper,
ToolCompressionMetrics,
ToolMetricsCollector,
get_tool_metrics,
reset_tool_metrics,
wrap_tools_with_headroom,
)
__all__ = [
"HeadroomToolWrapper",
"ToolCompressionMetrics",
"ToolMetricsCollector",
"wrap_tools_with_headroom",
"get_tool_metrics",
"reset_tool_metrics",
]

361
headroom/integrations/crewai/agents.py Normal file
View file

@ -0,0 +1,361 @@
"""CrewAI agent tool integration with output compression.
This module provides HeadroomToolWrapper and wrap_tools_with_headroom
for wrapping CrewAI tools to automatically compress their outputs
and track per-tool compression metrics.
Mirrors the LangChain agent integration pattern but targets CrewAI's
BaseTool interface (BaseTool.run -> _run -> result -> agent context).
Example:
from crewai.tools.base_tool import tool
from headroom.integrations.crewai import wrap_tools_with_headroom
@tool
def search_database(query: str) -> str:
\"\"\"Search the database.\"\"\"
return json.dumps({"results": [...], "total": 1000})
wrapped = wrap_tools_with_headroom(
[search_database],
min_chars_to_compress=1000,
)
"""
from __future__ import annotations
import logging
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any
try:
from crewai.tools.base_tool import BaseTool
CREWAI_AVAILABLE = True
except ImportError:
CREWAI_AVAILABLE = False
BaseTool = object # type: ignore[misc,assignment]
from headroom.integrations.mcp import compress_tool_result
logger = logging.getLogger(__name__)
def _check_crewai_available() -> None:
"""Raise ImportError if CrewAI is not installed."""
if not CREWAI_AVAILABLE:
raise ImportError(
"CrewAI is required for this integration. Install with: pip install crewai"
)
@dataclass
class ToolCompressionMetrics:
"""Metrics from a single tool compression.
Attributes:
tool_name: Name of the tool that was invoked.
timestamp: When the compression occurred.
chars_before: Character count of the original output.
chars_after: Character count after compression.
chars_saved: Characters removed by compression.
compression_ratio: Ratio of compressed to original size.
was_compressed: Whether compression was actually applied.
"""
tool_name: str
timestamp: datetime
chars_before: int
chars_after: int
chars_saved: int
compression_ratio: float
was_compressed: bool
@dataclass
class ToolMetricsCollector:
"""Collects compression metrics across all tool invocations.
Attributes:
metrics: List of per-invocation metrics.
"""
metrics: list[ToolCompressionMetrics] = field(default_factory=list)
def add(self, metric: ToolCompressionMetrics) -> None:
"""Add a metric entry.
Args:
metric: The compression metrics to record.
"""
self.metrics.append(metric)
if len(self.metrics) > 1000:
self.metrics = self.metrics[-1000:]
def get_summary(self) -> dict[str, Any]:
"""Get summary statistics.
Returns:
Dict with total_invocations, total_compressions,
total_chars_saved, average_compression_ratio, and
per-tool breakdown.
"""
if not self.metrics:
return {
"total_invocations": 0,
"total_compressions": 0,
"total_chars_saved": 0,
}
compressed = [m for m in self.metrics if m.was_compressed]
return {
"total_invocations": len(self.metrics),
"total_compressions": len(compressed),
"total_chars_saved": sum(m.chars_saved for m in self.metrics),
"average_compression_ratio": (
sum(m.compression_ratio for m in compressed) / len(compressed) if compressed else 0
),
"by_tool": self._get_by_tool_stats(),
}
def _get_by_tool_stats(self) -> dict[str, dict[str, Any]]:
"""Get per-tool statistics."""
by_tool: dict[str, list[ToolCompressionMetrics]] = {}
for m in self.metrics:
if m.tool_name not in by_tool:
by_tool[m.tool_name] = []
by_tool[m.tool_name].append(m)
result = {}
for name, tool_metrics in by_tool.items():
compressed = [m for m in tool_metrics if m.was_compressed]
result[name] = {
"invocations": len(tool_metrics),
"compressions": len(compressed),
"chars_saved": sum(m.chars_saved for m in tool_metrics),
}
return result
# Global metrics collector
_global_metrics = ToolMetricsCollector()
def get_tool_metrics() -> ToolMetricsCollector:
"""Get the global tool metrics collector.
Returns:
The global ToolMetricsCollector instance.
"""
return _global_metrics
def reset_tool_metrics() -> None:
"""Reset global tool metrics."""
global _global_metrics
_global_metrics = ToolMetricsCollector()
class HeadroomToolWrapper(BaseTool): # type: ignore[misc]
"""Wraps a CrewAI BaseTool to compress its output.
Applies Headroom compression to tool outputs, particularly useful
for tools that return large JSON arrays, search results, database
query results, or verbose log output.
The wrapper preserves the original tool's name, description, and
argument schema so it can be used as a drop-in replacement.
Example:
from crewai.tools.base_tool import tool
from headroom.integrations.crewai import HeadroomToolWrapper
@tool
def search(query: str) -> str:
\"\"\"Search and return results.\"\"\"
return json.dumps({"results": [...]})
wrapped = HeadroomToolWrapper(search)
result = wrapped.run(query="python tutorials")
Attributes:
name: Tool name (inherited from wrapped tool).
description: Tool description (inherited from wrapped tool).
"""
name: str = ""
description: str = ""
_inner: BaseTool
_min_chars: int
_metrics: ToolMetricsCollector
def __init__(
self,
tool: BaseTool,
min_chars_to_compress: int = 1000,
metrics_collector: ToolMetricsCollector | None = None,
) -> None:
"""Initialize HeadroomToolWrapper.
Args:
tool: The CrewAI BaseTool to wrap.
min_chars_to_compress: Minimum character count for output
before compression is applied. Default 1000.
metrics_collector: Collector for metrics. Uses global
collector if not specified.
"""
_check_crewai_available()
original_description = tool.description
super().__init__(
name=tool.name,
description=tool.description,
args_schema=tool.args_schema,
result_schema=getattr(tool, "result_schema", None),
cache_function=tool.cache_function,
result_as_answer=tool.result_as_answer,
max_usage_count=tool.max_usage_count,
)
# CrewAI's BaseTool rewrites description during construction
# (appends schema text). Restore the original.
self.description = original_description
self._inner = tool
self._min_chars = min_chars_to_compress
self._metrics = metrics_collector or _global_metrics
def _run(self, *args: Any, **kwargs: Any) -> Any:
"""Execute the wrapped tool and compress output.
Args:
*args: Positional arguments for the tool.
**kwargs: Keyword arguments for the tool.
Returns:
Compressed tool output as string.
"""
raw = self._inner.run(*args, **kwargs)
return self._compress_and_record(str(raw))
async def _arun(self, *args: Any, **kwargs: Any) -> Any:
"""Execute the wrapped tool asynchronously and compress output.
Args:
*args: Positional arguments for the tool.
**kwargs: Keyword arguments for the tool.
Returns:
Compressed tool output as string.
"""
raw = await self._inner.arun(*args, **kwargs)
return self._compress_and_record(str(raw))
def _compress_and_record(self, output: str) -> str:
"""Compress output and record metrics.
Args:
output: Tool output string.
Returns:
Compressed output, or original if below threshold or on error.
"""
chars_before = len(output)
if chars_before < self._min_chars:
self._record_metrics(output, output, was_compressed=False)
return output
try:
compressed = compress_tool_result(
content=output,
tool_name=self.name,
)
except Exception as e:
logger.debug("Tool compression failed for %s: %s", self.name, e)
self._record_metrics(output, output, was_compressed=False)
return output
self._record_metrics(output, compressed, was_compressed=True)
return compressed
def _record_metrics(self, original: str, compressed: str, was_compressed: bool) -> None:
"""Record compression metrics.
Args:
original: Original output.
compressed: Compressed output.
was_compressed: Whether compression was applied.
"""
chars_before = len(original)
chars_after = len(compressed)
chars_saved = chars_before - chars_after
metric = ToolCompressionMetrics(
tool_name=self.name,
timestamp=datetime.now(),
chars_before=chars_before,
chars_after=chars_after,
chars_saved=max(0, chars_saved),
compression_ratio=chars_after / chars_before if chars_before > 0 else 1.0,
was_compressed=was_compressed and chars_saved > 0,
)
self._metrics.add(metric)
if was_compressed and chars_saved > 0:
logger.info(
"HeadroomToolWrapper[%s]: %d -> %d chars (%d saved, %.1f%% of original)",
self.name,
chars_before,
chars_after,
chars_saved,
metric.compression_ratio * 100,
)
def wrap_tools_with_headroom(
tools: list[BaseTool],
min_chars_to_compress: int = 1000,
metrics_collector: ToolMetricsCollector | None = None,
) -> list[BaseTool]:
"""Wrap multiple CrewAI tools with Headroom compression.
Convenience function to wrap all tools in a list at once.
Each wrapped tool preserves the original's name, description,
and argument schema.
Args:
tools: List of CrewAI tools to wrap.
min_chars_to_compress: Minimum output size for compression.
metrics_collector: Shared metrics collector for all tools.
Returns:
List of wrapped tools.
Example:
from crewai import Agent
from crewai.tools.base_tool import tool
from headroom.integrations.crewai import wrap_tools_with_headroom
@tool
def search(query: str) -> str:
\"\"\"Search the database.\"\"\"
return json.dumps(results)
wrapped = wrap_tools_with_headroom([search])
agent = Agent(role="Researcher", tools=wrapped, ...)
"""
_check_crewai_available()
collector = metrics_collector or _global_metrics
return [
HeadroomToolWrapper(
tool=t,
min_chars_to_compress=min_chars_to_compress,
metrics_collector=collector,
)
for t in tools
]

8
pyproject.toml
View file

@ -210,6 +210,14 @@ agno = [
strands = [
"strands-agents>=0.1.0",
]
# CrewAI agent framework integration
crewai = [
"crewai>=1.0",
]
# AutoGen agent framework integration
autogen = [
"autogen-agentchat>=0.7",
]
# MCP server for Claude Code integration
mcp = [
"mcp>=1.0.0",

0
tests/test_integrations/autogen/__init__.py Normal file
View file

241
tests/test_integrations/autogen/test_agents.py Normal file
View file

@ -0,0 +1,241 @@
"""Tests for AutoGen agent tool integration.
Tests cover:
1. ToolCompressionMetrics - Dataclass for tool compression metrics
2. ToolMetricsCollector - Collector for compression metrics
3. HeadroomToolWrapper - Wrapper for AutoGen FunctionTool with compression
4. wrap_tools_with_headroom - Convenience function for wrapping multiple tools
5. get_tool_metrics / reset_tool_metrics - Global metrics access
"""
import asyncio
import json
from datetime import datetime
from unittest.mock import patch
import pytest
try:
from autogen_core import CancellationToken
from autogen_core.tools import FunctionTool
AUTOGEN_AVAILABLE = True
except ImportError:
AUTOGEN_AVAILABLE = False
pytestmark = pytest.mark.skipif(not AUTOGEN_AVAILABLE, reason="AutoGen not installed")
def _make_large_output(n: int = 200) -> str:
"""Create a large JSON string to trigger compression."""
return json.dumps({"items": [{"id": i, "data": "x" * 50} for i in range(n)]})
def _run_async(coro):
"""Helper to run async code in tests."""
return asyncio.run(coro)
# Sample tool functions
def big_lookup(query: str) -> str:
"""Look up data and return a large result."""
return _make_large_output()
def small_lookup(query: str) -> str:
"""Look up data and return a small result."""
return "ok"
async def async_lookup(query: str) -> str:
"""Async tool that returns a large result."""
return _make_large_output()
class TestToolCompressionMetrics:
"""Tests for ToolCompressionMetrics dataclass."""
def test_create_metrics(self):
from headroom.integrations.autogen.agents import ToolCompressionMetrics
metrics = ToolCompressionMetrics(
tool_name="search",
timestamp=datetime.now(),
chars_before=5000,
chars_after=2000,
chars_saved=3000,
compression_ratio=0.4,
was_compressed=True,
)
assert metrics.tool_name == "search"
assert metrics.chars_before == 5000
assert metrics.chars_saved == 3000
assert metrics.was_compressed is True
def test_metrics_all_fields_required(self):
from headroom.integrations.autogen.agents import ToolCompressionMetrics
with pytest.raises(TypeError):
ToolCompressionMetrics() # type: ignore[call-arg]
class TestToolMetricsCollector:
"""Tests for ToolMetricsCollector."""
def test_empty_summary(self):
from headroom.integrations.autogen.agents import ToolMetricsCollector
collector = ToolMetricsCollector()
summary = collector.get_summary()
assert summary["total_invocations"] == 0
assert summary["total_compressions"] == 0
def test_add_and_summary(self):
from headroom.integrations.autogen.agents import (
ToolCompressionMetrics,
ToolMetricsCollector,
)
collector = ToolMetricsCollector()
collector.add(
ToolCompressionMetrics(
tool_name="search",
timestamp=datetime.now(),
chars_before=5000,
chars_after=2000,
chars_saved=3000,
compression_ratio=0.4,
was_compressed=True,
)
)
summary = collector.get_summary()
assert summary["total_invocations"] == 1
assert summary["total_compressions"] == 1
assert summary["total_chars_saved"] == 3000
assert "search" in summary["by_tool"]
class TestGlobalMetrics:
"""Tests for global metrics functions."""
def test_get_and_reset(self):
from headroom.integrations.autogen.agents import get_tool_metrics, reset_tool_metrics
metrics = get_tool_metrics()
assert metrics is not None
reset_tool_metrics()
assert get_tool_metrics() is not metrics
class TestHeadroomToolWrapper:
"""Tests for HeadroomToolWrapper."""
@patch("headroom.integrations.autogen.agents.compress_tool_result")
def test_skips_short_output(self, mock_compress):
from headroom.integrations.autogen.agents import HeadroomToolWrapper, ToolMetricsCollector
tool = FunctionTool(small_lookup, description="Small", name="small_lookup")
collector = ToolMetricsCollector()
wrapper = HeadroomToolWrapper(tool, min_chars_to_compress=1000, metrics_collector=collector)
result = _run_async(wrapper.wrapped_tool.run_json({"query": "test"}, CancellationToken()))
assert str(result) == "ok"
mock_compress.assert_not_called()
assert collector.get_summary()["total_compressions"] == 0
@patch("headroom.integrations.autogen.agents.compress_tool_result")
def test_compresses_large_output(self, mock_compress):
from headroom.integrations.autogen.agents import HeadroomToolWrapper, ToolMetricsCollector
mock_compress.return_value = "compressed"
tool = FunctionTool(big_lookup, description="Big", name="big_lookup")
collector = ToolMetricsCollector()
wrapper = HeadroomToolWrapper(tool, min_chars_to_compress=100, metrics_collector=collector)
result = _run_async(wrapper.wrapped_tool.run_json({"query": "test"}, CancellationToken()))
assert str(result) == "compressed"
mock_compress.assert_called_once()
assert collector.get_summary()["total_compressions"] == 1
@patch(
"headroom.integrations.autogen.agents.compress_tool_result",
side_effect=RuntimeError("boom"),
)
def test_passes_through_on_error(self, mock_compress):
from headroom.integrations.autogen.agents import HeadroomToolWrapper, ToolMetricsCollector
large = _make_large_output()
tool = FunctionTool(big_lookup, description="Big", name="big_lookup")
collector = ToolMetricsCollector()
wrapper = HeadroomToolWrapper(tool, min_chars_to_compress=100, metrics_collector=collector)
result = _run_async(wrapper.wrapped_tool.run_json({"query": "test"}, CancellationToken()))
assert str(result) == large
assert collector.get_summary()["total_compressions"] == 0
def test_preserves_tool_metadata(self):
from headroom.integrations.autogen.agents import HeadroomToolWrapper
tool = FunctionTool(big_lookup, description="Look up data", name="big_lookup")
wrapper = HeadroomToolWrapper(tool)
assert wrapper.name == "big_lookup"
assert wrapper.description == "Look up data"
assert wrapper.wrapped_tool.name == "big_lookup"
@patch("headroom.integrations.autogen.agents.compress_tool_result")
def test_wraps_async_tool(self, mock_compress):
from headroom.integrations.autogen.agents import HeadroomToolWrapper, ToolMetricsCollector
mock_compress.return_value = "compressed"
tool = FunctionTool(async_lookup, description="Async", name="async_lookup")
collector = ToolMetricsCollector()
wrapper = HeadroomToolWrapper(tool, min_chars_to_compress=100, metrics_collector=collector)
result = _run_async(wrapper.wrapped_tool.run_json({"query": "test"}, CancellationToken()))
assert str(result) == "compressed"
assert collector.get_summary()["total_compressions"] == 1
class TestWrapToolsWithHeadroom:
"""Tests for wrap_tools_with_headroom convenience function."""
@patch("headroom.integrations.autogen.agents.compress_tool_result")
def test_wraps_multiple_tools(self, mock_compress):
from headroom.integrations.autogen.agents import wrap_tools_with_headroom
tool1 = FunctionTool(big_lookup, description="Big", name="big_lookup")
tool2 = FunctionTool(small_lookup, description="Small", name="small_lookup")
wrapped = wrap_tools_with_headroom([tool1, tool2])
assert len(wrapped) == 2
assert wrapped[0].name == "big_lookup"
assert wrapped[1].name == "small_lookup"
@patch("headroom.integrations.autogen.agents.compress_tool_result")
def test_shared_metrics(self, mock_compress):
from headroom.integrations.autogen.agents import (
ToolMetricsCollector,
wrap_tools_with_headroom,
)
mock_compress.return_value = "compressed"
tool1 = FunctionTool(big_lookup, description="Big", name="big_lookup")
tool2 = FunctionTool(big_lookup, description="Big2", name="big_lookup_2")
collector = ToolMetricsCollector()
wrapped = wrap_tools_with_headroom(
[tool1, tool2],
min_chars_to_compress=100,
metrics_collector=collector,
)
_run_async(wrapped[0].run_json({"query": "a"}, CancellationToken()))
_run_async(wrapped[1].run_json({"query": "b"}, CancellationToken()))
summary = collector.get_summary()
assert summary["total_invocations"] == 2
assert summary["total_compressions"] == 2

0
tests/test_integrations/crewai/__init__.py Normal file
View file

262
tests/test_integrations/crewai/test_agents.py Normal file
View file

@ -0,0 +1,262 @@
"""Tests for CrewAI agent tool integration.
Tests cover:
1. ToolCompressionMetrics - Dataclass for tool compression metrics
2. ToolMetricsCollector - Collector for compression metrics
3. HeadroomToolWrapper - Wrapper for CrewAI tools with compression
4. wrap_tools_with_headroom - Convenience function for wrapping multiple tools
5. get_tool_metrics / reset_tool_metrics - Global metrics access
"""
from datetime import datetime
from unittest.mock import patch
import pytest
try:
from crewai.tools.base_tool import tool as crewai_tool
CREWAI_AVAILABLE = True
except ImportError:
CREWAI_AVAILABLE = False
pytestmark = pytest.mark.skipif(not CREWAI_AVAILABLE, reason="CrewAI not installed")
def _make_large_output(n: int = 200) -> str:
"""Create a large JSON string to trigger compression."""
import json
return json.dumps({"items": [{"id": i, "data": "x" * 50} for i in range(n)]})
class TestToolCompressionMetrics:
"""Tests for ToolCompressionMetrics dataclass."""
def test_create_metrics(self):
from headroom.integrations.crewai.agents import ToolCompressionMetrics
metrics = ToolCompressionMetrics(
tool_name="search",
timestamp=datetime.now(),
chars_before=5000,
chars_after=2000,
chars_saved=3000,
compression_ratio=0.4,
was_compressed=True,
)
assert metrics.tool_name == "search"
assert metrics.chars_before == 5000
assert metrics.chars_saved == 3000
assert metrics.was_compressed is True
def test_metrics_all_fields_required(self):
from headroom.integrations.crewai.agents import ToolCompressionMetrics
with pytest.raises(TypeError):
ToolCompressionMetrics() # type: ignore[call-arg]
class TestToolMetricsCollector:
"""Tests for ToolMetricsCollector."""
def test_empty_summary(self):
from headroom.integrations.crewai.agents import ToolMetricsCollector
collector = ToolMetricsCollector()
summary = collector.get_summary()
assert summary["total_invocations"] == 0
assert summary["total_compressions"] == 0
def test_add_and_summary(self):
from headroom.integrations.crewai.agents import (
ToolCompressionMetrics,
ToolMetricsCollector,
)
collector = ToolMetricsCollector()
collector.add(
ToolCompressionMetrics(
tool_name="search",
timestamp=datetime.now(),
chars_before=5000,
chars_after=2000,
chars_saved=3000,
compression_ratio=0.4,
was_compressed=True,
)
)
summary = collector.get_summary()
assert summary["total_invocations"] == 1
assert summary["total_compressions"] == 1
assert summary["total_chars_saved"] == 3000
assert "search" in summary["by_tool"]
def test_caps_at_1000(self):
from headroom.integrations.crewai.agents import (
ToolCompressionMetrics,
ToolMetricsCollector,
)
collector = ToolMetricsCollector()
for _i in range(1050):
collector.add(
ToolCompressionMetrics(
tool_name="t",
timestamp=datetime.now(),
chars_before=100,
chars_after=100,
chars_saved=0,
compression_ratio=1.0,
was_compressed=False,
)
)
assert len(collector.metrics) == 1000
class TestGlobalMetrics:
"""Tests for global metrics functions."""
def test_get_and_reset(self):
from headroom.integrations.crewai.agents import get_tool_metrics, reset_tool_metrics
metrics = get_tool_metrics()
assert metrics is not None
reset_tool_metrics()
assert get_tool_metrics() is not metrics
class TestHeadroomToolWrapper:
"""Tests for HeadroomToolWrapper."""
@patch("headroom.integrations.crewai.agents.compress_tool_result")
def test_skips_short_output(self, mock_compress):
from headroom.integrations.crewai.agents import HeadroomToolWrapper, ToolMetricsCollector
@crewai_tool
def small_tool(query: str) -> str:
"""Return small output."""
return "short"
collector = ToolMetricsCollector()
wrapper = HeadroomToolWrapper(
small_tool,
min_chars_to_compress=1000,
metrics_collector=collector,
)
result = wrapper.run(query="test")
assert result == "short"
mock_compress.assert_not_called()
assert collector.get_summary()["total_compressions"] == 0
@patch("headroom.integrations.crewai.agents.compress_tool_result")
def test_compresses_large_output(self, mock_compress):
from headroom.integrations.crewai.agents import HeadroomToolWrapper, ToolMetricsCollector
large = _make_large_output()
mock_compress.return_value = "compressed"
@crewai_tool
def big_tool(query: str) -> str:
"""Return large output."""
return large
collector = ToolMetricsCollector()
wrapper = HeadroomToolWrapper(
big_tool,
min_chars_to_compress=100,
metrics_collector=collector,
)
result = wrapper.run(query="test")
assert result == "compressed"
mock_compress.assert_called_once()
assert collector.get_summary()["total_compressions"] == 1
@patch(
"headroom.integrations.crewai.agents.compress_tool_result",
side_effect=RuntimeError("boom"),
)
def test_passes_through_on_error(self, mock_compress):
from headroom.integrations.crewai.agents import HeadroomToolWrapper, ToolMetricsCollector
large = _make_large_output()
@crewai_tool
def flaky_tool(query: str) -> str:
"""Return large output."""
return large
collector = ToolMetricsCollector()
wrapper = HeadroomToolWrapper(
flaky_tool,
min_chars_to_compress=100,
metrics_collector=collector,
)
result = wrapper.run(query="test")
assert result == large
assert collector.get_summary()["total_compressions"] == 0
def test_preserves_tool_metadata(self):
from headroom.integrations.crewai.agents import HeadroomToolWrapper
@crewai_tool
def my_fn(x: int) -> str:
"""Do something useful."""
return str(x)
wrapper = HeadroomToolWrapper(my_fn)
assert wrapper.name == "my_fn"
assert wrapper.description == "Do something useful."
class TestWrapToolsWithHeadroom:
"""Tests for wrap_tools_with_headroom convenience function."""
@patch("headroom.integrations.crewai.agents.compress_tool_result")
def test_wraps_multiple_tools(self, mock_compress):
from headroom.integrations.crewai.agents import wrap_tools_with_headroom
@crewai_tool
def tool_a(q: str) -> str:
"""Tool A."""
return "a"
@crewai_tool
def tool_b(q: str) -> str:
"""Tool B."""
return "b"
wrapped = wrap_tools_with_headroom([tool_a, tool_b])
assert len(wrapped) == 2
assert wrapped[0].name == "tool_a"
assert wrapped[1].name == "tool_b"
@patch("headroom.integrations.crewai.agents.compress_tool_result")
def test_shared_metrics(self, mock_compress):
from headroom.integrations.crewai.agents import (
ToolMetricsCollector,
wrap_tools_with_headroom,
)
large = _make_large_output()
mock_compress.return_value = "compressed"
@crewai_tool
def big(q: str) -> str:
"""Big tool."""
return large
collector = ToolMetricsCollector()
wrapped = wrap_tools_with_headroom(
[big],
min_chars_to_compress=100,
metrics_collector=collector,
)
wrapped[0].run(q="test")
assert collector.get_summary()["total_invocations"] == 1