diff --git a/headroom/backends/base.py b/headroom/backends/base.py index 8bead4742..9aa77204e 100644 --- a/headroom/backends/base.py +++ b/headroom/backends/base.py @@ -117,6 +117,29 @@ class Backend(ABC): """ ... + async def send_openai_message( + self, + body: dict[str, Any], + headers: dict[str, str], + ) -> BackendResponse: + """Send an OpenAI-format message request. + + Unlike send_message(), this takes OpenAI-format input and returns + OpenAI-format output (no Anthropic conversion). Optional - only + implemented by backends that support OpenAI-compatible APIs. + + Args: + body: Request body in OpenAI chat completion format. + headers: Request headers. + + Returns: + BackendResponse with body in OpenAI chat completion format. + + Raises: + NotImplementedError: If backend doesn't support OpenAI format. + """ + raise NotImplementedError(f"{self.name} backend does not support OpenAI format") + async def close(self) -> None: # noqa: B027 """Clean up resources (e.g., close HTTP clients).""" pass diff --git a/headroom/backends/litellm.py b/headroom/backends/litellm.py index ebc577216..880712dca 100644 --- a/headroom/backends/litellm.py +++ b/headroom/backends/litellm.py @@ -198,6 +198,15 @@ PROVIDER_REGISTRY: dict[str, ProviderConfig] = { uses_region=True, env_vars=["AZURE_API_KEY", "AZURE_API_BASE"], ), + "databricks": ProviderConfig( + name="databricks", + display_name="Databricks", + model_map={}, # Pass through - Databricks uses custom model names + pass_through=True, + uses_region=False, + env_vars=["DATABRICKS_API_KEY", "DATABRICKS_API_BASE"], + model_format_hint="databricks-meta-llama-3-1-70b-instruct, databricks-dbrx-instruct, etc.", + ), } @@ -599,3 +608,134 @@ class LiteLLMBackend(Backend): async def close(self) -> None: # noqa: B027 """Clean up (no-op for LiteLLM).""" pass + + async def send_openai_message( + self, + body: dict[str, Any], + headers: dict[str, str], + ) -> BackendResponse: + """Send OpenAI-format message via LiteLLM. + + Unlike send_message(), this takes OpenAI-format input and returns + OpenAI-format output (no Anthropic conversion). + + Args: + body: OpenAI chat completion request body + headers: Request headers (ignored, auth from env vars) + + Returns: + BackendResponse with OpenAI-format body + """ + original_model = body.get("model", "gpt-4") + litellm_model = self.map_model_id(original_model) + + try: + # Build kwargs - messages already in OpenAI format + kwargs: dict[str, Any] = { + "model": litellm_model, + "messages": body.get("messages", []), + } + + # Pass through OpenAI parameters + for param in [ + "max_tokens", + "temperature", + "top_p", + "stop", + "tools", + "tool_choice", + "response_format", + "seed", + "n", + ]: + if param in body: + kwargs[param] = body[param] + + # Provider-specific config + if self.provider == "bedrock" and self.region: + kwargs["aws_region_name"] = self.region + elif self.provider == "databricks": + # Databricks uses env vars for auth + pass + + logger.debug(f"LiteLLM OpenAI request: model={litellm_model}") + + # Make the call + response = await acompletion(**kwargs) + + # Convert ModelResponse to dict (OpenAI format) + response_dict = { + "id": response.id, + "object": "chat.completion", + "created": response.created, + "model": original_model, + "choices": [ + { + "index": c.index, + "message": { + "role": c.message.role, + "content": c.message.content, + **( + { + "tool_calls": [ + { + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments, + }, + } + for tc in c.message.tool_calls + ] + } + if c.message.tool_calls + else {} + ), + }, + "finish_reason": c.finish_reason, + } + for c in response.choices + ], + "usage": { + "prompt_tokens": response.usage.prompt_tokens, + "completion_tokens": response.usage.completion_tokens, + "total_tokens": response.usage.total_tokens, + }, + } + + return BackendResponse( + body=response_dict, + status_code=200, + headers={"content-type": "application/json"}, + ) + + except Exception as e: + logger.error(f"LiteLLM OpenAI error: {e}") + + # Map to OpenAI error format + error_type = "api_error" + status_code = 500 + + error_str = str(e).lower() + if "authentication" in error_str or "credentials" in error_str: + error_type = "invalid_api_key" + status_code = 401 + elif "rate" in error_str or "limit" in error_str: + error_type = "rate_limit_exceeded" + status_code = 429 + elif "not found" in error_str: + error_type = "model_not_found" + status_code = 404 + + return BackendResponse( + body={ + "error": { + "message": str(e), + "type": error_type, + "code": error_type, + } + }, + status_code=status_code, + error=str(e), + ) diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index e6eb5a0e6..037cd4cea 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -4015,6 +4015,60 @@ class HeadroomProxy: body["messages"] = optimized_messages if tools is not None: body["tools"] = tools + + # Route through LiteLLM backend if configured (Databricks, Bedrock, etc.) + if self.anthropic_backend is not None: + try: + # Use the backend's OpenAI-format method + backend_response = await self.anthropic_backend.send_openai_message(body, headers) + + if backend_response.error: + return JSONResponse( + status_code=backend_response.status_code, + content=backend_response.body, + ) + + # Track metrics + total_latency = (time.time() - start_time) * 1000 + usage = backend_response.body.get("usage", {}) + output_tokens = usage.get("completion_tokens", 0) + total_input_tokens = usage.get("prompt_tokens", optimized_tokens) + + await self.metrics.record_request( + provider=self.anthropic_backend.name, + model=model, + input_tokens=total_input_tokens, + output_tokens=output_tokens, + tokens_saved=tokens_saved, + latency_ms=total_latency, + cached=False, + overhead_ms=optimization_latency, + ) + + if tokens_saved > 0: + logger.info( + f"[{request_id}] {model}: {original_tokens:,} → {optimized_tokens:,} " + f"(saved {tokens_saved:,} tokens) via {self.anthropic_backend.name}" + ) + + return JSONResponse( + status_code=backend_response.status_code, + content=backend_response.body, + ) + except Exception as e: + logger.error(f"[{request_id}] Backend error: {e}") + return JSONResponse( + status_code=500, + content={ + "error": { + "message": str(e), + "type": "api_error", + "code": "backend_error", + } + }, + ) + + # Direct OpenAI API (no backend configured) url = f"{self.OPENAI_API_URL}/v1/chat/completions" try: @@ -4196,6 +4250,59 @@ class HeadroomProxy: headers=response_headers, ) + # ========================================================================= + # Databricks Native API + # ========================================================================= + + async def handle_databricks_invocations( + self, + request: Request, + model: str, + ) -> Response | StreamingResponse: + """Handle Databricks native /serving-endpoints/{model}/invocations endpoint. + + This enables using the Databricks CLI directly with Headroom: + databricks serving-endpoints query --profile HEADROOM --json '{"messages": [...]}' + + The request/response format is identical to OpenAI chat completions, + so we inject the model from the path and delegate to handle_openai_chat. + """ + request_id = await self._next_request_id() + + try: + body = await request.json() + except Exception as e: + logger.error(f"[{request_id}] Failed to parse Databricks request body: {e}") + return JSONResponse( + status_code=400, + content={ + "error": {"message": f"Invalid JSON: {e}", "type": "invalid_request_error"} + }, + ) + + # Inject model from path into body (Databricks CLI passes model in URL, not body) + body["model"] = model + + logger.info(f"[{request_id}] Databricks invocation: model={model}") + + # Create a new request with the modified body + # We reuse the OpenAI chat handler since the format is identical + from starlette.requests import Request as StarletteRequest + + # Build new scope with the body already parsed + scope = dict(request.scope) + + # Create a simple receive function that returns our modified body + body_bytes = json.dumps(body).encode() + + async def receive(): + return {"type": "http.request", "body": body_bytes} + + modified_request = StarletteRequest(scope, receive) + + # Delegate to the OpenAI chat handler (same format) + return await self.handle_openai_chat(modified_request) + # ========================================================================= # OpenAI Batch API with Compression # ========================================================================= @@ -6186,6 +6293,21 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: """Gemini countTokens API with compression applied.""" return await proxy.handle_gemini_count_tokens(request, model) + # ========================================================================= + # Databricks Native Endpoints + # ========================================================================= + + @app.post("/serving-endpoints/{model}/invocations") + async def databricks_invocations(request: Request, model: str): + """Databricks native serving endpoint - compatible with Databricks CLI. + + This allows using the Databricks CLI directly with Headroom proxy: + databricks serving-endpoints query --profile HEADROOM --json '{"messages": [...]}' + + The request format is identical to OpenAI chat completions. + """ + return await proxy.handle_databricks_invocations(request, model) + # ========================================================================= # Passthrough Endpoints (no compression needed) # ========================================================================= diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 3ec3960d6..61a8ad7d9 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -45,10 +45,45 @@ from typing import Any from ..config import DEFAULT_EXCLUDE_TOOLS, TransformResult from ..tokenizer import Tokenizer from .base import Transform -from .content_detector import ContentType, detect_content_type +from .content_detector import ContentType, DetectionResult, detect_content_type logger = logging.getLogger(__name__) +# Use Magika-based detector if available, fallback to regex +_magika_detector = None +_USE_MAGIKA = False +try: + from ..compression.detector import get_detector + + _magika_detector = get_detector(prefer_magika=True) + _USE_MAGIKA = True + logger.info("ContentRouter: Using Magika ML-based content detection") +except ImportError: + logger.debug("Magika not available, using regex-based detection") + + +def _detect_content(content: str) -> DetectionResult: + """Detect content type using Magika if available, else regex fallback.""" + if _USE_MAGIKA and _magika_detector: + result = _magika_detector.detect(content) + # Map Magika ContentType to router's expected format + type_map = { + "json": ContentType.JSON_ARRAY, + "code": ContentType.SOURCE_CODE, + "log": ContentType.BUILD_OUTPUT, + "markdown": ContentType.PLAIN_TEXT, + "text": ContentType.PLAIN_TEXT, + "unknown": ContentType.PLAIN_TEXT, + } + mapped_type = type_map.get(result.content_type.value, ContentType.PLAIN_TEXT) + return DetectionResult( + content_type=mapped_type, + confidence=result.confidence, + metadata={"language": result.language, "raw_label": result.raw_label}, + ) + else: + return detect_content_type(content) + def _create_content_signature( content_type: str, @@ -595,7 +630,7 @@ class ContentRouter(Transform): return CompressionStrategy.MIXED # 2. Detect content type from content itself - detection = detect_content_type(content) + detection = _detect_content(content) return self._strategy_from_detection(detection) def _strategy_from_detection(self, detection: Any) -> CompressionStrategy: @@ -1183,7 +1218,7 @@ class ContentRouter(Transform): continue # Detect content type for protection decisions - detection = detect_content_type(content) + detection = _detect_content(content) is_code = detection.content_type == ContentType.SOURCE_CODE # Protection 2: Don't compress recent CODE