mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
efd2ac1ca4
`.gitattributes` declares `*.py text eol=lf` and `*.sh text eol=lf`, but 74 files (73 .py, 1 .sh) are stored in the index with CRLF line endings, violating that contract. Every macOS/Linux clone reports these files as "modified" on fresh checkout because git's diff engine sees the stored bytes don't match the attribute contract, even though the working tree and index match byte-for-byte. Running `git add --renormalize .` rewrites each affected blob so the stored form matches the attribute declaration. No semantic changes — every affected file's diff is "N insertions, N deletions" with inserts and deletes being the same lines modulo line endings. Follow-up commit adds `.git-blame-ignore-revs` so `git blame` / GitHub blame skip this mechanical commit. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
274 lines
9.5 KiB
Python
274 lines
9.5 KiB
Python
"""Tests for the one-function compress() API and integrations."""
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
from headroom.compress import CompressResult, compress
|
|
from headroom.hooks import CompressionHooks
|
|
|
|
try:
|
|
from starlette.applications import Starlette
|
|
from starlette.requests import Request
|
|
from starlette.responses import JSONResponse
|
|
from starlette.routing import Route
|
|
from starlette.testclient import TestClient
|
|
|
|
from headroom.integrations.asgi import CompressionMiddleware
|
|
|
|
HAS_STARLETTE = True
|
|
except ImportError:
|
|
HAS_STARLETTE = False
|
|
|
|
|
|
# =============================================================================
|
|
# Tests: compress() function
|
|
# =============================================================================
|
|
|
|
|
|
class TestCompressFunction:
|
|
def test_empty_messages(self):
|
|
result = compress([], model="test")
|
|
assert result.messages == []
|
|
assert result.tokens_saved == 0
|
|
|
|
def test_small_messages_passthrough(self):
|
|
"""Small messages below compression threshold pass through unchanged."""
|
|
messages = [{"role": "user", "content": "hello"}]
|
|
result = compress(messages, model="gpt-4o")
|
|
assert result.messages[0]["content"] == "hello"
|
|
assert result.tokens_saved == 0
|
|
|
|
def test_returns_compress_result(self):
|
|
result = compress([{"role": "user", "content": "hi"}])
|
|
assert isinstance(result, CompressResult)
|
|
assert hasattr(result, "messages")
|
|
assert hasattr(result, "tokens_saved")
|
|
assert hasattr(result, "compression_ratio")
|
|
assert hasattr(result, "transforms_applied")
|
|
|
|
def test_large_tool_output_compressed(self):
|
|
"""Large JSON tool output should be compressed."""
|
|
big_data = json.dumps(
|
|
[
|
|
{"id": i, "status": "active", "name": f"item_{i}", "value": i * 17}
|
|
for i in range(200)
|
|
]
|
|
)
|
|
messages = [
|
|
{"role": "user", "content": "What are the top items?"},
|
|
{"role": "tool", "content": big_data, "tool_call_id": "call_1"},
|
|
]
|
|
result = compress(messages, model="gpt-4o")
|
|
assert result.tokens_after <= result.tokens_before
|
|
assert len(result.messages) == 2
|
|
|
|
def test_compact_json_counts_tokens_not_whitespace(self):
|
|
"""Compact JSON arrays should still compress under token thresholds."""
|
|
numbers = [42.0 + i * 0.1 for i in range(200)]
|
|
messages = [
|
|
{"role": "system", "content": "You are helpful."},
|
|
{"role": "user", "content": "Show metrics"},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {"name": "get_metrics", "arguments": "{}"},
|
|
}
|
|
],
|
|
},
|
|
{"role": "tool", "tool_call_id": "call_1", "content": json.dumps(numbers)},
|
|
]
|
|
|
|
result = compress(messages, min_tokens_to_compress=250)
|
|
|
|
assert result.tokens_saved > 0
|
|
assert any(
|
|
transform.startswith("router:smart_crusher") for transform in result.transforms_applied
|
|
)
|
|
|
|
def test_optimize_false_passthrough(self):
|
|
"""optimize=False returns messages unchanged."""
|
|
messages = [{"role": "user", "content": "hello world " * 100}]
|
|
result = compress(messages, optimize=False)
|
|
assert result.messages is messages
|
|
assert result.tokens_saved == 0
|
|
|
|
def test_with_custom_hooks(self):
|
|
"""Hooks are called when provided."""
|
|
calls = []
|
|
|
|
class TrackingHooks(CompressionHooks):
|
|
def pre_compress(self, messages, ctx):
|
|
calls.append(("pre", len(messages)))
|
|
return messages
|
|
|
|
def compute_biases(self, messages, ctx):
|
|
calls.append(("biases", len(messages)))
|
|
return {}
|
|
|
|
def post_compress(self, event):
|
|
calls.append(("post", event.tokens_saved))
|
|
|
|
big_data = json.dumps([{"id": i, "status": "active"} for i in range(100)])
|
|
messages = [
|
|
{"role": "user", "content": "analyze"},
|
|
{"role": "tool", "content": big_data, "tool_call_id": "c1"},
|
|
]
|
|
compress(messages, hooks=TrackingHooks())
|
|
|
|
assert any(c[0] == "pre" for c in calls)
|
|
assert any(c[0] == "biases" for c in calls)
|
|
|
|
|
|
class TestCompressResultFields:
|
|
def test_fields_populated(self):
|
|
big_data = json.dumps([{"id": i, "type": "log"} for i in range(100)])
|
|
messages = [
|
|
{"role": "user", "content": "summarize"},
|
|
{"role": "tool", "content": big_data, "tool_call_id": "c1"},
|
|
]
|
|
result = compress(messages, model="claude-sonnet-4-5-20250929")
|
|
assert result.tokens_before > 0
|
|
assert result.tokens_after >= 0
|
|
assert result.tokens_saved >= 0
|
|
assert 0.0 <= result.compression_ratio <= 1.0
|
|
|
|
|
|
# =============================================================================
|
|
# Tests: ASGI CompressionMiddleware (requires starlette)
|
|
# =============================================================================
|
|
|
|
|
|
def _make_asgi_app(middleware_kwargs=None):
|
|
"""Create a test ASGI app with CompressionMiddleware."""
|
|
|
|
async def chat_endpoint(request: Request) -> JSONResponse:
|
|
body = await request.json()
|
|
return JSONResponse(
|
|
{
|
|
"model": "gpt-4o",
|
|
"choices": [{"message": {"content": "response"}}],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 5},
|
|
"_message_count": len(body.get("messages", [])),
|
|
}
|
|
)
|
|
|
|
async def health(request: Request) -> JSONResponse:
|
|
return JSONResponse({"status": "ok"})
|
|
|
|
app = Starlette(
|
|
routes=[
|
|
Route("/health", health),
|
|
Route("/v1/chat/completions", chat_endpoint, methods=["POST"]),
|
|
Route("/v1/messages", chat_endpoint, methods=["POST"]),
|
|
]
|
|
)
|
|
app.add_middleware(CompressionMiddleware, **(middleware_kwargs or {}))
|
|
return app
|
|
|
|
|
|
@pytest.mark.skipif(not HAS_STARLETTE, reason="starlette not installed")
|
|
class TestASGIMiddleware:
|
|
def test_non_llm_paths_passthrough(self):
|
|
app = _make_asgi_app()
|
|
client = TestClient(app)
|
|
resp = client.get("/health")
|
|
assert resp.status_code == 200
|
|
assert resp.json()["status"] == "ok"
|
|
|
|
def test_small_messages_passthrough(self):
|
|
app = _make_asgi_app()
|
|
client = TestClient(app)
|
|
resp = client.post(
|
|
"/v1/chat/completions",
|
|
json={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]},
|
|
)
|
|
assert resp.status_code == 200
|
|
|
|
def test_large_messages_compressed(self):
|
|
"""Large tool output should be compressed by middleware."""
|
|
app = _make_asgi_app()
|
|
client = TestClient(app)
|
|
|
|
big_data = json.dumps([{"id": i, "status": "active"} for i in range(200)])
|
|
resp = client.post(
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": "gpt-4o",
|
|
"messages": [
|
|
{"role": "user", "content": "analyze"},
|
|
{"role": "tool", "content": big_data, "tool_call_id": "c1"},
|
|
],
|
|
},
|
|
)
|
|
assert resp.status_code == 200
|
|
|
|
def test_anthropic_path(self):
|
|
"""Works with Anthropic /v1/messages path."""
|
|
app = _make_asgi_app()
|
|
client = TestClient(app)
|
|
resp = client.post(
|
|
"/v1/messages",
|
|
json={
|
|
"model": "claude-sonnet-4-5-20250929",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
},
|
|
)
|
|
assert resp.status_code == 200
|
|
|
|
def test_get_requests_passthrough(self):
|
|
"""GET requests to LLM paths pass through."""
|
|
app = _make_asgi_app()
|
|
client = TestClient(app)
|
|
resp = client.get("/v1/chat/completions")
|
|
assert resp.status_code in (200, 405)
|
|
|
|
|
|
# =============================================================================
|
|
# Tests: LiteLLM Callback
|
|
# =============================================================================
|
|
|
|
|
|
class TestLiteLLMCallback:
|
|
def test_callback_imports(self):
|
|
"""Verify the callback can be imported."""
|
|
from headroom.integrations.litellm_callback import HeadroomCallback
|
|
|
|
callback = HeadroomCallback()
|
|
assert callback.total_tokens_saved == 0
|
|
|
|
def test_callback_compresses_messages(self):
|
|
"""Callback compresses messages in pre_call_hook."""
|
|
import asyncio
|
|
|
|
from headroom.integrations.litellm_callback import HeadroomCallback
|
|
|
|
callback = HeadroomCallback()
|
|
|
|
big_data = json.dumps([{"id": i, "status": "active"} for i in range(200)])
|
|
data = {
|
|
"model": "gpt-4o",
|
|
"messages": [
|
|
{"role": "user", "content": "analyze"},
|
|
{"role": "tool", "content": big_data, "tool_call_id": "c1"},
|
|
],
|
|
}
|
|
|
|
result = asyncio.run(callback.async_pre_call_hook("key", data, "completion"))
|
|
assert result is data
|
|
|
|
def test_callback_ignores_non_completion(self):
|
|
"""Non-completion calls are passed through."""
|
|
import asyncio
|
|
|
|
from headroom.integrations.litellm_callback import HeadroomCallback
|
|
|
|
callback = HeadroomCallback()
|
|
data = {"messages": [{"role": "user", "content": "hi"}]}
|
|
|
|
result = asyncio.run(callback.async_pre_call_hook("key", data, "embedding"))
|
|
assert result is data
|