diff --git a/headroom/integrations/asgi.py b/headroom/integrations/asgi.py index bc9f2027d..8503363c7 100644 --- a/headroom/integrations/asgi.py +++ b/headroom/integrations/asgi.py @@ -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 diff --git a/headroom/integrations/litellm_callback.py b/headroom/integrations/litellm_callback.py index 413fc973a..50bf8d9ea 100644 --- a/headroom/integrations/litellm_callback.py +++ b/headroom/integrations/litellm_callback.py @@ -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: diff --git a/uv.lock b/uv.lock index 07b58ebf1..ff66741c3 100644 --- a/uv.lock +++ b/uv.lock @@ -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 },