diff --git a/CHANGELOG.md b/CHANGELOG.md index 4780ed3f3..44d033136 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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)) diff --git a/docs/content/docs/autogen.mdx b/docs/content/docs/autogen.mdx new file mode 100644 index 000000000..1495040b1 --- /dev/null +++ b/docs/content/docs/autogen.mdx @@ -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. diff --git a/docs/content/docs/crewai.mdx b/docs/content/docs/crewai.mdx new file mode 100644 index 000000000..ec4876b6e --- /dev/null +++ b/docs/content/docs/crewai.mdx @@ -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. diff --git a/headroom/integrations/__init__.py b/headroom/integrations/__init__.py index 3f50ef01f..a639960a4 100644 --- a/headroom/integrations/__init__.py +++ b/headroom/integrations/__init__.py @@ -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", ] diff --git a/headroom/integrations/autogen/__init__.py b/headroom/integrations/autogen/__init__.py new file mode 100644 index 000000000..132671a1c --- /dev/null +++ b/headroom/integrations/autogen/__init__.py @@ -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", +] diff --git a/headroom/integrations/autogen/agents.py b/headroom/integrations/autogen/agents.py new file mode 100644 index 000000000..a83abbd1a --- /dev/null +++ b/headroom/integrations/autogen/agents.py @@ -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 + ] diff --git a/headroom/integrations/crewai/__init__.py b/headroom/integrations/crewai/__init__.py new file mode 100644 index 000000000..5efb7a864 --- /dev/null +++ b/headroom/integrations/crewai/__init__.py @@ -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", +] diff --git a/headroom/integrations/crewai/agents.py b/headroom/integrations/crewai/agents.py new file mode 100644 index 000000000..e96de992a --- /dev/null +++ b/headroom/integrations/crewai/agents.py @@ -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 + ] diff --git a/pyproject.toml b/pyproject.toml index 174e60c96..4a4915a76 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/tests/test_integrations/autogen/__init__.py b/tests/test_integrations/autogen/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/test_integrations/autogen/test_agents.py b/tests/test_integrations/autogen/test_agents.py new file mode 100644 index 000000000..fb3049e4f --- /dev/null +++ b/tests/test_integrations/autogen/test_agents.py @@ -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 diff --git a/tests/test_integrations/crewai/__init__.py b/tests/test_integrations/crewai/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/test_integrations/crewai/test_agents.py b/tests/test_integrations/crewai/test_agents.py new file mode 100644 index 000000000..35c958c9d --- /dev/null +++ b/tests/test_integrations/crewai/test_agents.py @@ -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