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 LiteLLM backend routing for OpenAI endpoint and Magika content detection

Browse source 
- Route /v1/chat/completions through configured LiteLLM backend (Bedrock, Azure,
  Databricks, etc.) instead of hardcoding to OpenAI API
- Add send_openai_message() method to Backend base class and LiteLLMBackend
- Add Databricks provider to PROVIDER_REGISTRY
- Add /serving-endpoints/{model}/invocations endpoint for Databricks CLI compatibility
- Integrate Magika ML-based content detection in ContentRouter for improved accuracy
- Fall back to regex-based detection when Magika is unavailable
This commit is contained in:
chopratejas 2026-02-04 10:14:34 -08:00
parent 5e3aa3880f
commit c90cf39826
4 changed files with 323 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

23
headroom/backends/base.py
View file

@ -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

140
headroom/backends/litellm.py
View file

@ -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),
)

122
headroom/proxy/server.py
View file

@ -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 <model> --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 <model> --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)
# =========================================================================

41
headroom/transforms/content_router.py
View file

@ -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