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

Add Cloud mode to ASGI middleware and LiteLLM callback

Browse source 
Both CompressionMiddleware and HeadroomCallback now support a cloud mode
(api_key="hdr_xxx") that calls Headroom Cloud API for managed compression
with org-scoped CCR, TOIN learning, and analytics. Falls back to
HEADROOM_API_KEY env var. Local mode (default) is unchanged.

Also adds x-headroom-tokens-before/after response headers and updates
uv.lock with mcp extra and version bump to 0.3.3.
This commit is contained in:
chopratejas 2026-02-27 20:10:07 -08:00
parent 876949e638
commit 40369762dd
3 changed files with 225 additions and 62 deletions
Download patch file Download diff file Expand all files Collapse all files

142
headroom/integrations/asgi.py
View file

@ -3,41 +3,40 @@
Drop-in middleware for FastAPI, Starlette, LiteLLM proxy, or any ASGI app.
Intercepts LLM requests, compresses messages, forwards the smaller payload.
Local mode (compression runs in-process):
from headroom.integrations.asgi import CompressionMiddleware
app.add_middleware(CompressionMiddleware)
Cloud mode (managed CCR, TOIN, analytics via Headroom Cloud):
app.add_middleware(CompressionMiddleware, api_key="hdr_xxx")
Usage with LiteLLM proxy:
from litellm.proxy.proxy_server import app
from headroom.integrations.asgi import CompressionMiddleware
app.add_middleware(CompressionMiddleware)
app.add_middleware(CompressionMiddleware) # local
# OR
app.add_middleware(CompressionMiddleware, api_key="hdr_xxx") # cloud
Usage with any FastAPI app:
from fastapi import FastAPI
from headroom.integrations.asgi import CompressionMiddleware
app = FastAPI()
app.add_middleware(CompressionMiddleware)
# Your existing routes...
Configuration:
app.add_middleware(
CompressionMiddleware,
min_tokens=500, # Only compress if messages > 500 tokens
model_limit=200000, # Context window size
)
Cloud mode requires httpx: pip install httpx
"""
from __future__ import annotations
import json
import logging
import os
from typing import Any
from starlette.types import ASGIApp, Receive, Scope, Send
logger = logging.getLogger(__name__)
_DEFAULT_CLOUD_URL = "https://api.headroomlabs.ai"
# Paths that contain LLM messages to compress
_LLM_PATHS = (
"/v1/messages", # Anthropic
@ -50,13 +49,16 @@ _LLM_PATHS = (
class CompressionMiddleware:
"""ASGI middleware that compresses LLM request messages.
Intercepts POST requests to LLM endpoints, compresses the messages
using Headroom's full pipeline, and forwards the smaller payload.
Two modes:
- Local (default): Compresses in-process using headroom.compress().
- Cloud (api_key set): Calls Headroom Cloud API for managed compression
with org-scoped CCR, TOIN learning, and analytics dashboards.
Response headers include compression metrics:
- x-headroom-tokens-before: original token count
- x-headroom-tokens-after: compressed token count
- x-headroom-tokens-saved: tokens removed
- x-headroom-compressed: "true" if compression occurred
"""
def __init__(
@ -65,12 +67,26 @@ class CompressionMiddleware:
min_tokens: int = 500,
model_limit: int = 200000,
hooks: Any = None,
api_key: str | None = None,
api_url: str | None = None,
) -> None:
self.app = app
self._min_tokens = min_tokens
self._model_limit = model_limit
self._hooks = hooks
# Cloud mode: if api_key is set, compress via Headroom Cloud API
self._api_key = api_key or os.environ.get("HEADROOM_API_KEY", "").strip() or None
self._api_url = (
api_url or os.environ.get("HEADROOM_API_URL", "").strip() or _DEFAULT_CLOUD_URL
).rstrip("/")
self._client = None # Lazy-initialized httpx.AsyncClient
@property
def cloud_mode(self) -> bool:
"""Whether cloud compression is enabled."""
return self._api_key is not None
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] != "http":
await self.app(scope, receive, send)
@ -106,32 +122,33 @@ class CompressionMiddleware:
# Parse and compress
tokens_saved = 0
tokens_before = 0
tokens_after = 0
try:
body_json = json.loads(full_body)
messages = body_json.get("messages", [])
model = body_json.get("model", "")
if messages:
from headroom.compress import compress
if self._api_key:
result = await self._cloud_compress(messages, model)
else:
result = self._local_compress(messages, model)
result = compress(
messages=messages,
model=model or "claude-sonnet-4-5-20250929",
model_limit=self._model_limit,
hooks=self._hooks,
)
if result.tokens_saved > 0:
body_json["messages"] = result.messages
if result and result.get("tokens_saved", 0) > 0:
body_json["messages"] = result["messages"]
full_body = json.dumps(body_json).encode("utf-8")
tokens_saved = result.tokens_saved
tokens_saved = result["tokens_saved"]
tokens_before = result.get("tokens_before", 0)
tokens_after = result.get("tokens_after", 0)
logger.info(
"Headroom: %d→%d tokens (saved %d, %.0f%%)",
result.tokens_before,
result.tokens_after,
result.tokens_saved,
result.compression_ratio * 100,
"Headroom%s: %d→%d tokens (saved %d, %.0f%%)",
" Cloud" if self._api_key else "",
tokens_before,
tokens_after,
tokens_saved,
result.get("compression_ratio", 0) * 100,
)
except (json.JSONDecodeError, TypeError, KeyError) as e:
@ -153,8 +170,63 @@ class CompressionMiddleware:
if message["type"] == "http.response.start" and tokens_saved > 0:
headers = list(message.get("headers", []))
headers.append((b"x-headroom-compressed", b"true"))
headers.append((b"x-headroom-tokens-before", str(tokens_before).encode()))
headers.append((b"x-headroom-tokens-after", str(tokens_after).encode()))
headers.append((b"x-headroom-tokens-saved", str(tokens_saved).encode()))
message = {**message, "headers": headers}
await send(message)
await self.app(scope, modified_receive, metrics_send)
def _local_compress(self, messages: list[dict], model: str) -> dict[str, Any] | None:
"""Compress locally using headroom.compress()."""
from headroom.compress import compress
result = compress(
messages=messages,
model=model or "claude-sonnet-4-5-20250929",
model_limit=self._model_limit,
hooks=self._hooks,
)
return {
"messages": result.messages,
"tokens_before": result.tokens_before,
"tokens_after": result.tokens_after,
"tokens_saved": result.tokens_saved,
"compression_ratio": result.compression_ratio,
}
async def _cloud_compress(self, messages: list[dict], model: str) -> dict[str, Any] | None:
"""Compress via Headroom Cloud API (managed CCR, TOIN, analytics)."""
if self._client is None:
try:
import httpx
except ImportError as e:
raise ImportError(
"httpx is required for Headroom Cloud mode: pip install httpx"
) from e
self._client = httpx.AsyncClient(timeout=30.0)
client = self._client
assert client is not None
resp = await client.post(
f"{self._api_url}/v1/saas/compress",
headers={
"X-Headroom-Key": self._api_key,
"Content-Type": "application/json",
},
content=json.dumps(
{
"messages": messages,
"model": model or "claude-sonnet-4-5-20250929",
"model_limit": self._model_limit,
}
),
)
if resp.status_code != 200:
logger.warning("Headroom Cloud API error: %d %s", resp.status_code, resp.text[:200])
return None
result: dict[str, Any] = resp.json()
return result

130
headroom/integrations/litellm_callback.py
View file

@ -1,38 +1,51 @@
"""LiteLLM callback — add Headroom compression to LiteLLM with one line.
# Local mode (compression runs in-process):
import litellm
from headroom.integrations.litellm_callback import HeadroomCallback
litellm.callbacks = [HeadroomCallback()]
# All LiteLLM calls now get compressed automatically.
# Or with custom config:
litellm.callbacks = [HeadroomCallback(min_tokens=1000)]
# Cloud mode (managed CCR, TOIN, analytics via Headroom Cloud):
litellm.callbacks = [HeadroomCallback(api_key="hdr_xxx")]
Works with LiteLLM's completion(), acompletion(), and proxy modes.
Cloud mode requires httpx: pip install httpx
"""
from __future__ import annotations
import json
import logging
from typing import Any
logger = logging.getLogger(__name__)
_DEFAULT_CLOUD_URL = "https://api.headroomlabs.ai"
class HeadroomCallback:
"""LiteLLM callback that compresses messages before each API call.
Implements LiteLLM's CustomLogger interface (async_pre_call_hook).
Compresses messages using Headroom's full pipeline, reducing token
usage across all providers LiteLLM supports.
Usage:
import litellm
from headroom.integrations.litellm_callback import HeadroomCallback
Two modes:
- Local (default): Compresses in-process using headroom.compress().
- Cloud (api_key set): Calls Headroom Cloud API for managed compression
with org-scoped CCR, TOIN learning, and analytics dashboards.
Usage (local):
litellm.callbacks = [HeadroomCallback()]
response = litellm.completion(model="gpt-4o", messages=[...])
Usage (cloud):
litellm.callbacks = [HeadroomCallback(api_key="hdr_xxx")]
Usage (cloud with LiteLLM proxy config):
# litellm_config.yaml
litellm_settings:
callbacks: [headroom.integrations.litellm_callback.HeadroomCallback]
environment_variables:
HEADROOM_API_KEY: "hdr_xxx"
"""
def __init__(
@ -40,17 +53,34 @@ class HeadroomCallback:
min_tokens: int = 500,
model_limit: int = 200000,
hooks: Any = None,
api_key: str | None = None,
api_url: str | None = None,
) -> None:
self._min_tokens = min_tokens
self._model_limit = model_limit
self._hooks = hooks
self._total_saved = 0
# Cloud mode: if api_key is set, compress via Headroom Cloud API
# Falls back to HEADROOM_API_KEY env var
import os
self._api_key = api_key or os.environ.get("HEADROOM_API_KEY", "").strip() or None
self._api_url = (
api_url or os.environ.get("HEADROOM_API_URL", "").strip() or _DEFAULT_CLOUD_URL
).rstrip("/")
self._client = None # Lazy-initialized httpx.AsyncClient
@property
def total_tokens_saved(self) -> int:
"""Total tokens saved across all calls."""
return self._total_saved
@property
def cloud_mode(self) -> bool:
"""Whether cloud compression is enabled."""
return self._api_key is not None
async def async_pre_call_hook(
self,
user_api_key: str,
@ -68,24 +98,21 @@ class HeadroomCallback:
return data
try:
from headroom.compress import compress
if self._api_key:
result = await self._cloud_compress(messages, model)
else:
result = self._local_compress(messages, model)
result = compress(
messages=messages,
model=model or "claude-sonnet-4-5-20250929",
model_limit=self._model_limit,
hooks=self._hooks,
)
if result.tokens_saved > 0:
data["messages"] = result.messages
self._total_saved += result.tokens_saved
if result and result.get("tokens_saved", 0) > 0:
data["messages"] = result["messages"]
self._total_saved += result["tokens_saved"]
logger.info(
"Headroom: %d→%d tokens (saved %d, %.0f%%) [total saved: %d]",
result.tokens_before,
result.tokens_after,
result.tokens_saved,
result.compression_ratio * 100,
"Headroom%s: %d→%d tokens (saved %d, %.0f%%) [total saved: %d]",
" Cloud" if self._api_key else "",
result["tokens_before"],
result["tokens_after"],
result["tokens_saved"],
result.get("compression_ratio", 0) * 100,
self._total_saved,
)
@ -94,6 +121,59 @@ class HeadroomCallback:
return data
def _local_compress(self, messages: list[dict], model: str) -> dict[str, Any] | None:
"""Compress locally using headroom.compress()."""
from headroom.compress import compress
result = compress(
messages=messages,
model=model or "claude-sonnet-4-5-20250929",
model_limit=self._model_limit,
hooks=self._hooks,
)
return {
"messages": result.messages,
"tokens_before": result.tokens_before,
"tokens_after": result.tokens_after,
"tokens_saved": result.tokens_saved,
"compression_ratio": result.compression_ratio,
}
async def _cloud_compress(self, messages: list[dict], model: str) -> dict[str, Any] | None:
"""Compress via Headroom Cloud API (managed CCR, TOIN, analytics)."""
if self._client is None:
try:
import httpx
except ImportError as e:
raise ImportError(
"httpx is required for Headroom Cloud mode: pip install httpx"
) from e
self._client = httpx.AsyncClient(timeout=30.0)
client = self._client
assert client is not None
resp = await client.post(
f"{self._api_url}/v1/saas/compress",
headers={
"X-Headroom-Key": self._api_key,
"Content-Type": "application/json",
},
content=json.dumps(
{
"messages": messages,
"model": model or "claude-sonnet-4-5-20250929",
"model_limit": self._model_limit,
}
),
)
if resp.status_code != 200:
logger.warning("Headroom Cloud API error: %d %s", resp.status_code, resp.text[:200])
return None
result: dict[str, Any] = resp.json()
return result
async def async_success_handler(
self, kwargs: dict, response: Any, start_time: Any, end_time: Any
) -> None:

15
uv.lock generated
View file

@ -1208,7 +1208,7 @@ wheels = [
[[package]]
name = "headroom-ai"
version = "0.3.0"
version = "0.3.3"
source = { editable = "." }
dependencies = [
{ name = "accelerate" },
@ -1240,6 +1240,7 @@ all = [
{ name = "jinja2" },
{ name = "llmlingua" },
{ name = "lm-eval" },
{ name = "mcp" },
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.netflix.net/simple" }, marker = "python_full_version < '3.11'" },
{ name = "numpy", version = "2.4.1", source = { registry = "https://pypi.netflix.net/simple" }, marker = "python_full_version >= '3.11'" },
{ name = "onnxruntime" },
@ -1296,6 +1297,10 @@ llmlingua = [
{ name = "torch" },
{ name = "transformers" },
]
mcp = [
{ name = "httpx" },
{ name = "mcp" },
]
memory = [
{ name = "hnswlib" },
{ name = "sqlite-vec" },
@ -1343,17 +1348,19 @@ requires-dist = [
{ name = "datasets", marker = "extra == 'evals'", specifier = ">=2.14.0" },
{ name = "datasets", marker = "extra == 'voice-train'", specifier = ">=2.14.0" },
{ name = "fastapi", marker = "extra == 'proxy'", specifier = ">=0.100.0" },
{ name = "headroom-ai", extras = ["relevance", "proxy", "reports", "llmlingua", "code", "evals", "memory", "voice", "html", "benchmark"], marker = "extra == 'all'" },
{ name = "headroom-ai", extras = ["relevance", "proxy", "reports", "llmlingua", "code", "evals", "memory", "voice", "html", "benchmark", "mcp"], marker = "extra == 'all'" },
{ name = "headroom-ai", extras = ["voice"], marker = "extra == 'voice-train'" },
{ name = "hnswlib", specifier = ">=0.8.0" },
{ name = "hnswlib", marker = "extra == 'dev'", specifier = ">=0.8.0" },
{ name = "hnswlib", marker = "extra == 'memory'", specifier = ">=0.8.0" },
{ name = "httpx", marker = "extra == 'mcp'", specifier = ">=0.24.0" },
{ name = "httpx", extras = ["http2"], marker = "extra == 'proxy'", specifier = ">=0.24.0" },
{ name = "jinja2", marker = "extra == 'reports'", specifier = ">=3.0.0" },
{ name = "langchain-ollama", marker = "extra == 'dev'", specifier = ">=0.2.0" },
{ name = "litellm", specifier = ">=1.0.0" },
{ name = "llmlingua", marker = "extra == 'llmlingua'", specifier = ">=0.2.0" },
{ name = "lm-eval", marker = "extra == 'benchmark'", specifier = ">=0.4.0" },
{ name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.0.0" },
{ name = "mypy", marker = "extra == 'dev'", specifier = ">=1.0.0" },
{ name = "numpy", marker = "extra == 'evals'", specifier = ">=1.24.0" },
{ name = "numpy", marker = "extra == 'relevance'", specifier = ">=1.24.0" },
@ -5076,6 +5083,10 @@ dependencies = [
{ name = "typing-extensions" },
]
wheels = [
{ url = "https://pypi.netflix.net/packages/19819580139/torch-2.10.0-2-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:2b980edd8d7c0a68c4e951ee1856334a43193f98730d97408fbd148c1a933313", size = 79411457 },
{ url = "https://pypi.netflix.net/packages/19819580140/torch-2.10.0-2-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:418997cb02d0a0f1497cf6a09f63166f9f5df9f3e16c8a716ab76a72127c714f", size = 79423467 },
{ url = "https://pypi.netflix.net/packages/19819580141/torch-2.10.0-2-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:13ec4add8c3faaed8d13e0574f5cd4a323c11655546f91fbe6afa77b57423574", size = 79498202 },
{ url = "https://pypi.netflix.net/packages/19819580142/torch-2.10.0-2-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:e521c9f030a3774ed770a9c011751fb47c4d12029a3d6522116e48431f2ff89e", size = 79498254 },
{ url = "https://pypi.netflix.net/packages/19646079874/torch-2.10.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:5276fa790a666ee8becaffff8acb711922252521b28fbce5db7db5cf9cb2026d", size = 145992962 },
{ url = "https://pypi.netflix.net/packages/19646073760/torch-2.10.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:aaf663927bcd490ae971469a624c322202a2a1e68936eb952535ca4cd3b90444", size = 915599237 },
{ url = "https://pypi.netflix.net/packages/19646080742/torch-2.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:a4be6a2a190b32ff5c8002a0977a25ea60e64f7ba46b1be37093c141d9c49aeb", size = 113720931 },