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 OpenRouter backend support with provider registry pattern

Browse source 
- Add native support for OpenRouter API via LiteLLM backend
- Introduce PROVIDER_REGISTRY pattern to eliminate scattered if/else blocks
- New providers can now be added with a single registry entry

Features:
- `headroom proxy --backend openrouter` routes requests to OpenRouter
- Pass-through model naming (anthropic/claude-3.5-sonnet, openai/gpt-4o, etc.)
- CLI shows provider-specific setup instructions from registry

Usage:
  export OPENROUTER_API_KEY="sk-or-v1-..."
  headroom proxy --backend openrouter

Also fixes mypy type errors in mcp_server.py
This commit is contained in:
chopratejas 2026-01-31 14:58:09 -08:00
parent bbbf0c5a7e
commit 8d1731b20c
5 changed files with 169 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

18
README.md
View file

@ -290,6 +290,24 @@ claude
**Do NOT include** `CLAUDE_CODE_USE_BEDROCK` - Headroom handles the Bedrock routing.
</details>
**Using OpenRouter?** Access 400+ models through a single API:
```bash
# OpenRouter - Terminal 1: Start proxy
export OPENROUTER_API_KEY="sk-or-v1-..."
headroom proxy --backend openrouter
# OpenRouter - Terminal 2: Run your client
export ANTHROPIC_API_KEY="sk-ant-dummy" # Any value works! Headroom ignores it.
export ANTHROPIC_BASE_URL="http://localhost:8787"
# Use OpenRouter model names in your requests:
# - anthropic/claude-3.5-sonnet
# - openai/gpt-4o
# - google/gemini-pro
# - meta-llama/llama-3-70b-instruct
# See all models: https://openrouter.ai/models
```
```bash
# Google Vertex AI
headroom proxy --backend vertex_ai --region us-central1

112
headroom/backends/litellm.py
View file

@ -4,6 +4,7 @@ Uses LiteLLM to support 100+ providers with minimal code:
- AWS Bedrock: model="bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0"
- Azure OpenAI: model="azure/gpt-4"
- Google Vertex: model="vertex_ai/claude-3-5-sonnet"
- OpenRouter: model="openrouter/anthropic/claude-3.5-sonnet"
- And many more...
LiteLLM handles all the auth and format translation internally.
@ -14,6 +15,7 @@ from __future__ import annotations
import logging
import uuid
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from typing import Any
from .base import Backend, BackendResponse, StreamEvent
@ -31,11 +33,26 @@ except ImportError:
acompletion = None # type: ignore
# Model mapping: Anthropic model IDs -> LiteLLM model strings
# IMPORTANT: Claude 4+ models require inference profiles (us.anthropic.* or global.anthropic.*)
# Direct model IDs (anthropic.*) don't support on-demand throughput for newer models.
# See: https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles.html
BEDROCK_MODEL_MAP = {
# =============================================================================
# Provider Registry - Add new providers here!
# =============================================================================
@dataclass
class ProviderConfig:
"""Configuration for a LiteLLM provider."""
name: str # Provider identifier (e.g., "bedrock", "openrouter")
display_name: str # Human-readable name (e.g., "AWS Bedrock", "OpenRouter")
model_map: dict[str, str] = field(default_factory=dict) # Anthropic -> provider model map
pass_through: bool = False # If True, prepend provider/ to any model
uses_region: bool = True # Whether region is relevant for this provider
env_vars: list[str] = field(default_factory=list) # Required env vars
model_format_hint: str = "" # Hint for model naming (shown in help)
# Model mappings for providers that need translation
_BEDROCK_MODEL_MAP = {
# Claude 4.5 (requires inference profiles)
"claude-opus-4-5-20251101": "bedrock/us.anthropic.claude-opus-4-5-20251101-v1:0",
"claude-sonnet-4-5-20250929": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
@ -57,7 +74,7 @@ BEDROCK_MODEL_MAP = {
"claude-3-haiku-20240307": "bedrock/us.anthropic.claude-3-haiku-20240307-v1:0",
}
VERTEX_MODEL_MAP = {
_VERTEX_MODEL_MAP = {
"claude-3-5-sonnet-20241022": "vertex_ai/claude-3-5-sonnet-v2@20241022",
"claude-3-5-sonnet-20240620": "vertex_ai/claude-3-5-sonnet@20240620",
"claude-3-opus-20240229": "vertex_ai/claude-3-opus@20240229",
@ -66,14 +83,65 @@ VERTEX_MODEL_MAP = {
}
# Provider Registry - to add a new provider, just add an entry here!
PROVIDER_REGISTRY: dict[str, ProviderConfig] = {
"bedrock": ProviderConfig(
name="bedrock",
display_name="AWS Bedrock",
model_map=_BEDROCK_MODEL_MAP,
uses_region=True,
env_vars=["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION"],
),
"vertex_ai": ProviderConfig(
name="vertex_ai",
display_name="Google Vertex AI",
model_map=_VERTEX_MODEL_MAP,
uses_region=True,
env_vars=["GOOGLE_APPLICATION_CREDENTIALS"],
),
"openrouter": ProviderConfig(
name="openrouter",
display_name="OpenRouter",
model_map={}, # No static map - pass through
pass_through=True,
uses_region=False,
env_vars=["OPENROUTER_API_KEY"],
model_format_hint="anthropic/claude-3.5-sonnet, openai/gpt-4o, etc.",
),
"azure": ProviderConfig(
name="azure",
display_name="Azure OpenAI",
model_map={},
uses_region=True,
env_vars=["AZURE_API_KEY", "AZURE_API_BASE"],
),
}
def get_provider_config(provider: str) -> ProviderConfig:
"""Get provider config, with fallback for unknown providers."""
if provider in PROVIDER_REGISTRY:
return PROVIDER_REGISTRY[provider]
# Fallback for unknown providers - basic pass-through
return ProviderConfig(
name=provider,
display_name=provider.upper(),
model_map={},
pass_through=True,
)
class LiteLLMBackend(Backend):
"""Backend using LiteLLM for multi-provider support.
Supports any provider LiteLLM supports:
- bedrock: AWS Bedrock (uses AWS credentials)
- vertex_ai: Google Vertex AI (uses GCP credentials)
- openrouter: OpenRouter (400+ models via single API)
- azure: Azure OpenAI (uses Azure credentials)
- And 100+ more...
To add a new provider, just add an entry to PROVIDER_REGISTRY above.
"""
def __init__(
@ -85,7 +153,7 @@ class LiteLLMBackend(Backend):
"""Initialize LiteLLM backend.
Args:
provider: LiteLLM provider prefix (bedrock, vertex_ai, azure, etc.)
provider: LiteLLM provider prefix (bedrock, vertex_ai, openrouter, etc.)
region: Cloud region (provider-specific)
**kwargs: Additional provider-specific config
"""
@ -98,16 +166,13 @@ class LiteLLMBackend(Backend):
self.region = region
self.kwargs = kwargs
# Select model map based on provider
if provider == "bedrock":
self._model_map = BEDROCK_MODEL_MAP
# Set AWS region for litellm
if region:
litellm.set_verbose = False # Reduce noise
elif provider == "vertex_ai":
self._model_map = VERTEX_MODEL_MAP
else:
self._model_map = {}
# Get provider config from registry
self._config = get_provider_config(provider)
self._model_map = self._config.model_map
# Provider-specific setup
if provider == "bedrock" and region:
litellm.set_verbose = False # Reduce noise
logger.info(f"LiteLLM backend initialized (provider={provider})")
@ -117,10 +182,18 @@ class LiteLLMBackend(Backend):
def map_model_id(self, anthropic_model: str) -> str:
"""Map Anthropic model ID to LiteLLM model string."""
# Check direct mapping
# Check direct mapping first
if anthropic_model in self._model_map:
return self._model_map[anthropic_model]
# Pass-through providers: prepend provider prefix
if self._config.pass_through:
# If already has provider prefix, use as-is
if anthropic_model.startswith(f"{self.provider}/"):
return anthropic_model
# Otherwise prepend provider/
return f"{self.provider}/{anthropic_model}"
# If already has provider prefix, use as-is
if "/" in anthropic_model:
return anthropic_model
@ -130,6 +203,9 @@ class LiteLLMBackend(Backend):
def supports_model(self, model: str) -> bool:
"""Check if model is supported."""
# Pass-through providers accept any model
if self._config.pass_through:
return True
return "claude" in model.lower() or model in self._model_map
def _convert_messages_for_litellm(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:

4
headroom/ccr/mcp_server.py
View file

@ -43,8 +43,8 @@ try:
MCP_AVAILABLE = True
except ImportError:
MCP_AVAILABLE = False
Server = None
stdio_server = None
Server = None # type: ignore[assignment,misc]
stdio_server = None # type: ignore[assignment]
# Try to import httpx for proxy communication
try:

56
headroom/cli/proxy.py
View file

@ -77,8 +77,8 @@ from .main import main
"--backend",
default="anthropic",
help=(
"API backend: 'anthropic' (direct), 'bedrock' (AWS), "
"or 'litellm-<provider>' (e.g., litellm-bedrock, litellm-vertex)"
"API backend: 'anthropic' (direct), 'bedrock' (AWS), 'openrouter' (OpenRouter), "
"or 'litellm-<provider>' (e.g., litellm-vertex)"
),
)
@click.option(
@ -185,13 +185,36 @@ def proxy(
effective_region = bedrock_region or region
backend_status = "Anthropic (direct API)"
backend_section = ""
if config.backend != "anthropic":
# Normalize: "bedrock" -> "litellm-bedrock"
backend_name = config.backend
if not backend_name.startswith("litellm-"):
backend_name = f"litellm-{backend_name}"
provider = backend_name.replace("litellm-", "")
backend_status = f"{provider.upper()} via LiteLLM (region={effective_region})"
# Get provider config from registry
from headroom.backends.litellm import get_provider_config
provider = config.backend.replace("litellm-", "")
provider_config = get_provider_config(provider)
# Build backend status
if provider_config.uses_region:
backend_status = (
f"{provider_config.display_name} via LiteLLM (region={effective_region})"
)
else:
backend_status = f"{provider_config.display_name} via LiteLLM"
# Build usage instructions from provider config
env_vars_str = (
", ".join(provider_config.env_vars) if provider_config.env_vars else "See docs"
)
backend_section = f"""
IMPORTANT for {provider_config.display_name} users:
1. Set credentials: {env_vars_str}
2. Set a dummy Anthropic key: ANTHROPIC_API_KEY="sk-ant-dummy"
(Headroom ignores this - it uses your {provider_config.display_name} credentials)
3. Set base URL: ANTHROPIC_BASE_URL=http://{config.host}:{config.port}"""
if provider_config.model_format_hint:
backend_section += f"\n 4. Use model names: {provider_config.model_format_hint}"
backend_section += "\n"
# Build memory section if enabled
memory_section = ""
@ -202,20 +225,7 @@ Memory:
- Tools: {"ENABLED" if config.memory_inject_tools else "DISABLED"} Context: {"ENABLED" if config.memory_inject_context else "DISABLED"}
"""
if config.memory_inject_tools:
memory_section += " - NOTE: Memory tools require ANTHROPIC_API_KEY (Claude Code subscription credentials have restrictions).\n"
# Build usage note for Bedrock
bedrock_section = ""
if config.backend != "anthropic":
provider = config.backend.replace("litellm-", "").upper()
bedrock_section = f"""
IMPORTANT for {provider} users (Claude Code / VS Code):
1. Set AWS credentials: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION
2. Set a dummy Anthropic key: ANTHROPIC_API_KEY="sk-ant-dummy"
(Headroom ignores this - it uses your AWS credentials for {provider})
3. Set base URL: ANTHROPIC_BASE_URL=http://{config.host}:{config.port}
4. Do NOT set: CLAUDE_CODE_USE_BEDROCK=1 (Headroom handles this)
"""
memory_section += " - NOTE: Memory tools require ANTHROPIC_API_KEY.\n"
click.echo(f"""
╔═══════════════════════════════════════════════════════════════════════╗
@ -231,7 +241,7 @@ Starting proxy server...
Caching: {"ENABLED" if config.cache_enabled else "DISABLED"}
Rate Limit: {"ENABLED" if config.rate_limit_enabled else "DISABLED"}
Memory: {memory_status}
{bedrock_section}
{backend_section}
Usage with Claude Code:
ANTHROPIC_BASE_URL=http://{config.host}:{config.port} claude

32
headroom/proxy/server.py
View file

@ -6259,16 +6259,20 @@ def run_server(
pool_info = f"max={config.max_connections}, keepalive={config.max_keepalive_connections}"
http2_status = "ENABLED" if config.http2 else "DISABLED"
# Backend status
# Backend status - use provider registry for display info
if config.backend == "anthropic":
backend_status = "ANTHROPIC (direct API)"
else:
# Normalize: "bedrock" -> "litellm-bedrock"
backend = config.backend
if not backend.startswith("litellm-"):
backend = f"litellm-{backend}"
provider = backend.replace("litellm-", "")
backend_status = f"{provider.upper()} via LiteLLM (region={config.bedrock_region})"
from headroom.backends.litellm import get_provider_config
provider = config.backend.replace("litellm-", "")
provider_config = get_provider_config(provider)
if provider_config.uses_region:
backend_status = (
f"{provider_config.display_name} via LiteLLM (region={config.bedrock_region})"
)
else:
backend_status = f"{provider_config.display_name} via LiteLLM"
print(f"""
╔══════════════════════════════════════════════════════════════════════╗
@ -6368,12 +6372,12 @@ if __name__ == "__main__":
"--openai-api-url", help=f"Custom OpenAI API URL (default: {HeadroomProxy.OPENAI_API_URL})"
)
# Backend (anthropic direct or bedrock)
# Backend (anthropic direct, bedrock, or openrouter)
parser.add_argument(
"--backend",
choices=["anthropic", "bedrock"],
choices=["anthropic", "bedrock", "openrouter"],
default="anthropic",
help="Backend for Anthropic API: 'anthropic' (direct) or 'bedrock' (AWS Bedrock)",
help="Backend for Anthropic API: 'anthropic' (direct), 'bedrock' (AWS), or 'openrouter' (OpenRouter)",
)
parser.add_argument(
"--bedrock-region",
@ -6384,6 +6388,10 @@ if __name__ == "__main__":
"--bedrock-profile",
help="AWS profile for Bedrock backend (default: use default credentials)",
)
parser.add_argument(
"--openrouter-api-key",
help="OpenRouter API key (or set OPENROUTER_API_KEY env var)",
)
# Connection pool (scalability)
parser.add_argument(
@ -6506,6 +6514,10 @@ if __name__ == "__main__":
cache_enabled = env_cache if not args.no_cache else False
rate_limit_enabled = env_rate_limit if not args.no_rate_limit else False
# Set OpenRouter API key from CLI if provided
if hasattr(args, "openrouter_api_key") and args.openrouter_api_key:
os.environ["OPENROUTER_API_KEY"] = args.openrouter_api_key
config = ProxyConfig(
host=_get_env_str("HEADROOM_HOST", args.host),
port=_get_env_int("HEADROOM_PORT", args.port),