"""LiteLLM-based backend for Headroom. 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. """ from __future__ import annotations import importlib.util import json import logging import os import uuid from collections.abc import AsyncIterator from dataclasses import dataclass, field from typing import Any from .base import Backend, BackendResponse, StreamEvent logger = logging.getLogger(__name__) # litellm calls `dotenv.load_dotenv()` during its own import, which loads # the project `.env` into `os.environ`. We don't want that side effect — # importing a backend module should not silently leak API keys into the # process. Snapshot `os.environ` around the import and undo any keys # litellm added. Same pattern as `headroom/pricing/litellm_pricing.py`. try: import os as _os _env_snapshot = set(_os.environ) import litellm from litellm import acompletion for _leaked_key in set(_os.environ) - _env_snapshot: del _os.environ[_leaked_key] del _env_snapshot, _os LITELLM_AVAILABLE = True except ImportError: LITELLM_AVAILABLE = False litellm = None # type: ignore acompletion = None # type: ignore # ============================================================================= # 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) # Cache for dynamically fetched inference profiles _bedrock_profiles_cache: dict[str, dict[str, str]] = {} # region -> model_map # Region prefix used in cross-region Bedrock inference profile IDs. # EU regions use "eu.", AP regions use "apac.", US (and everything else) use "us.". # ap-southeast-2 (Sydney/Australia) uses "au." — distinct from the rest of APAC. _BEDROCK_REGION_PREFIXES: dict[str, str] = { "eu": "eu", "ap-southeast-2": "au", "ap": "apac", } def _bedrock_region_prefix(region: str) -> str: """Return the inference-profile region prefix for an AWS region. AWS Bedrock cross-region inference profiles are prefixed with a geographic tag: ``us.``, ``eu.``, or ``apac.``. This helper maps an AWS region name (e.g. ``eu-west-1``) to the correct prefix. >>> _bedrock_region_prefix("us-east-1") 'us' >>> _bedrock_region_prefix("eu-central-1") 'eu' >>> _bedrock_region_prefix("ap-southeast-1") 'apac' """ for key, prefix in _BEDROCK_REGION_PREFIXES.items(): if region.startswith(key): return prefix return "us" def _build_bedrock_fallback_map(region: str) -> dict[str, str]: """Build a static Bedrock model map using the region prefix. When ``_fetch_bedrock_inference_profiles`` cannot reach the AWS API (wrong credentials, network error, permissions, etc.) we fall back to this map so that the proxy can still route requests. The map covers all currently GA Claude models on Bedrock. """ prefix = _bedrock_region_prefix(region) # Base model IDs without region prefix _CLAUDE_MODELS = [ # Claude 4.6 ("claude-opus-4-6", "anthropic.claude-opus-4-6-v1"), ("claude-sonnet-4-6", "anthropic.claude-sonnet-4-6"), # Claude 4.5 ("claude-sonnet-4-5-20250929", "anthropic.claude-sonnet-4-5-20250929-v1:0"), ("claude-opus-4-5-20251101", "anthropic.claude-opus-4-5-20251101-v1:0"), # Claude 4.1 ("claude-opus-4-1-20250805", "anthropic.claude-opus-4-1-20250805-v1:0"), # Claude 4 ("claude-sonnet-4-20250514", "anthropic.claude-sonnet-4-20250514-v1:0"), ("claude-opus-4-20250514", "anthropic.claude-opus-4-20250514-v1:0"), # Claude 3.7 ("claude-3-7-sonnet-20250219", "anthropic.claude-3-7-sonnet-20250219-v1:0"), # Claude 3.5 ("claude-3-5-sonnet-20241022", "anthropic.claude-3-5-sonnet-20241022-v2:0"), ("claude-3-5-sonnet-20240620", "anthropic.claude-3-5-sonnet-20240620-v1:0"), ("claude-3-5-haiku-20241022", "anthropic.claude-3-5-haiku-20241022-v1:0"), # Claude 3 ("claude-3-opus-20240229", "anthropic.claude-3-opus-20240229-v1:0"), ("claude-3-sonnet-20240229", "anthropic.claude-3-sonnet-20240229-v1:0"), ("claude-3-haiku-20240307", "anthropic.claude-3-haiku-20240307-v1:0"), # Haiku 4.5 ("claude-haiku-4-5-20251001", "anthropic.claude-haiku-4-5-20251001-v1:0"), ] return {name: f"bedrock/{prefix}.{model_id}" for name, model_id in _CLAUDE_MODELS} def _fetch_bedrock_inference_profiles( region: str | None, profile_name: str | None = None ) -> dict[str, str]: """Fetch available Bedrock inference profiles from AWS API. Uses boto3 list_inference_profiles() to get all available profiles for the given region, then builds a model map. If the API call fails (wrong credentials, network error, permission denied, etc.) the function logs a warning and returns a static fallback map so the proxy can still start. Args: region: AWS region (e.g., "us-east-1", "eu-central-1") profile_name: AWS named profile (e.g., "my-sso-profile"). When set, a boto3.Session is created with this profile name so the correct SSO or credential file is used. Falls back to ambient credentials (AWS_PROFILE env var, instance metadata, etc.) when not provided. Returns: Model map: anthropic_model_name -> bedrock inference profile ID """ region = region or "us-east-1" # Cache key includes profile_name so different profiles don't collide cache_key = f"{region}:{profile_name or ''}" if cache_key in _bedrock_profiles_cache: return _bedrock_profiles_cache[cache_key] model_map: dict[str, str] = {} try: import boto3 except ImportError: logger.warning( "boto3 is not installed — using static Bedrock model map. " "Install boto3 for dynamic model discovery: pip install boto3" ) model_map = _build_bedrock_fallback_map(region) _bedrock_profiles_cache[cache_key] = model_map return model_map try: session = boto3.Session(profile_name=profile_name) if profile_name else boto3.Session() bedrock_client = session.client("bedrock", region_name=region) response = bedrock_client.list_inference_profiles(typeEquals="SYSTEM_DEFINED") for profile in response.get("inferenceProfileSummaries", []): profile_id = profile.get("inferenceProfileId", "") # Only process Anthropic Claude profiles if "anthropic" not in profile_id.lower(): continue # Extract the standard model name from the profile ID # e.g., "us.anthropic.claude-sonnet-4-20250514-v1:0" -> "claude-sonnet-4-20250514" normalized = _normalize_bedrock_profile_id(profile_id) if normalized: model_map[normalized] = f"bedrock/{profile_id}" # Handle pagination if needed while response.get("nextToken"): response = bedrock_client.list_inference_profiles( typeEquals="SYSTEM_DEFINED", nextToken=response["nextToken"] ) for profile in response.get("inferenceProfileSummaries", []): profile_id = profile.get("inferenceProfileId", "") if "anthropic" not in profile_id.lower(): continue normalized = _normalize_bedrock_profile_id(profile_id) if normalized: model_map[normalized] = f"bedrock/{profile_id}" logger.info(f"Fetched {len(model_map)} Bedrock inference profiles for region {region}") except Exception as e: logger.warning( f"Failed to fetch Bedrock inference profiles for region {region}: {e}. " "Using static fallback model map." ) model_map = _build_bedrock_fallback_map(region) # Cache the result _bedrock_profiles_cache[cache_key] = model_map return model_map def _parse_bedrock_model_overrides(raw: str | None) -> dict[str, str]: """Parse the ``HEADROOM_BEDROCK_MODEL_MAP`` operator override. AWS discovery keys the model map by the *normalized model name*, so it cannot disambiguate application inference profiles that share one underlying model — e.g. a team where ``claude-sonnet-5-kenneth`` and ``claude-sonnet-5-jeremy`` both resolve to ``claude-sonnet-5``. When you need requests billed to a *specific* application profile (per-user cost attribution), pin the mapping explicitly here. The plain name Claude Code sends (kept plain so tool-search deferral stays on) resolves to your ARN. Format: comma-separated ``name=target`` pairs, where ``target`` is an application-inference-profile ARN (routed via the converse endpoint) or any LiteLLM model string. Whitespace around pairs is ignored; blank entries are skipped. HEADROOM_BEDROCK_MODEL_MAP="claude-sonnet-5=arn:aws:bedrock:...:application-inference-profile/x57j1esjrt66,claude-opus-4-8=arn:aws:bedrock:...:application-inference-profile/3dy9ytxuq2ci" """ overrides: dict[str, str] = {} if not raw: return overrides for pair in raw.split(","): pair = pair.strip() if not pair or "=" not in pair: continue name, _, target = pair.partition("=") name = name.strip() target = target.strip() if name and target: overrides[name] = target return overrides def _normalize_bedrock_profile_id(profile_id: str) -> str | None: """Extract standard Anthropic model name from Bedrock profile ID. Args: profile_id: e.g., "us.anthropic.claude-sonnet-4-20250514-v1:0" or "anthropic.claude-sonnet-4-20250514-v1:0" or "claude-sonnet-4-20250514" or "arn:aws:bedrock:...:application-inference-profile/..." Returns: Normalized name like "claude-sonnet-4-20250514", or None if not parseable """ import re # ARNs are opaque identifiers — cannot be normalized to a standard model name if profile_id.startswith("arn:aws:"): return None # Strip "bedrock/" prefix if present if profile_id.startswith("bedrock/"): profile_id = profile_id[8:] # Strip region prefix (us., eu., apac., au.) or the newer "global." # cross-region prefix used by current-gen profiles (e.g. # "global.anthropic.claude-sonnet-4-6"). for prefix in ["us.", "eu.", "apac.", "au.", "global."]: if profile_id.startswith(prefix): profile_id = profile_id[len(prefix) :] break # Strip "anthropic." prefix if profile_id.startswith("anthropic."): profile_id = profile_id[10:] # Must be a Claude model if not profile_id.startswith("claude"): return None # Strip version suffix. Legacy dated profiles use "-v1:0" / "-v2:0"; # newer undated profiles use a bare "-v1" (no colon/revision) or carry # no version suffix at all (e.g. "claude-opus-4-8"). Match all three # shapes so undated current-gen profiles normalize instead of # silently falling out of the resolvable model map. normalized = re.sub(r"-v\d+(?::\d+)?$", "", profile_id) return normalized if normalized else None # Legacy static map - kept for non-Bedrock providers _BEDROCK_MODEL_MAP: dict[str, str] = {} _VERTEX_MODEL_MAP = { # Claude 4.6 (latest, no date suffix) "claude-opus-4-6": "vertex_ai/claude-opus-4-6", "claude-sonnet-4-6": "vertex_ai/claude-sonnet-4-6", # Claude 4.5 "claude-sonnet-4-5-20250929": "vertex_ai/claude-sonnet-4-5@20250929", "claude-opus-4-5-20251101": "vertex_ai/claude-opus-4-5@20251101", # Claude 4.1 "claude-opus-4-1-20250805": "vertex_ai/claude-opus-4-1@20250805", # Claude 4 "claude-sonnet-4-20250514": "vertex_ai/claude-sonnet-4@20250514", "claude-opus-4-20250514": "vertex_ai/claude-opus-4@20250514", # Claude 3.7 "claude-3-7-sonnet-20250219": "vertex_ai/claude-3-7-sonnet@20250219", # Claude 3.5 "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-5-haiku-20241022": "vertex_ai/claude-3-5-haiku@20241022", # Claude 3 (haiku 3 deprecated, others retired) "claude-3-opus-20240229": "vertex_ai/claude-3-opus@20240229", "claude-3-sonnet-20240229": "vertex_ai/claude-3-sonnet@20240229", "claude-3-haiku-20240307": "vertex_ai/claude-3-haiku@20240307", # Haiku 4.5 "claude-haiku-4-5-20251001": "vertex_ai/claude-haiku-4-5@20251001", } # 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"], ), "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.", ), } 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, ) def _anthropic_usage_from_litellm(litellm_usage: Any) -> dict[str, Any]: """Map LiteLLM usage to Anthropic-shape usage, surfacing cache tokens. LiteLLM's ``prompt_tokens`` is the *total* prompt size including cached tokens, while Anthropic's ``input_tokens`` excludes tokens served from or written to the prompt cache. Without this mapping a working Bedrock prompt cache is invisible to non-streaming clients: they see the full prompt count and no cache fields, which looks exactly like the cache being broken (see #1345). The streaming/OpenAI paths already surface these fields. """ cache_read = int(getattr(litellm_usage, "cache_read_input_tokens", 0) or 0) cache_write = int(getattr(litellm_usage, "cache_creation_input_tokens", 0) or 0) details = getattr(litellm_usage, "prompt_tokens_details", None) if details is not None: cache_read = cache_read or int(getattr(details, "cached_tokens", 0) or 0) cache_write = cache_write or int(getattr(details, "cache_creation_tokens", 0) or 0) prompt_tokens = int(getattr(litellm_usage, "prompt_tokens", 0) or 0) usage: dict[str, Any] = { "input_tokens": max(prompt_tokens - cache_read - cache_write, 0), "output_tokens": getattr(litellm_usage, "completion_tokens", 0), } if cache_read or cache_write: usage["cache_read_input_tokens"] = cache_read usage["cache_creation_input_tokens"] = cache_write return usage def _convert_anthropic_tool(tool: dict[str, Any]) -> dict[str, Any]: """Convert Anthropic tool format to OpenAI function format. Anthropic: {"name": "...", "description": "...", "input_schema": {...}} OpenAI: {"type": "function", "function": {"name": "...", "description": "...", "parameters": {...}}} """ func: dict[str, Any] = {"name": tool.get("name", "")} if "description" in tool: func["description"] = tool["description"] if "input_schema" in tool: func["parameters"] = tool["input_schema"] return {"type": "function", "function": func} def _convert_tool_choice(choice: Any) -> Any: """Convert Anthropic tool_choice to OpenAI format. Anthropic: {"type": "auto"}, {"type": "any"}, {"type": "tool", "name": "..."} OpenAI: "auto", "required", {"type": "function", "function": {"name": "..."}} """ if isinstance(choice, str): return choice if isinstance(choice, dict): choice_type = choice.get("type", "auto") if choice_type == "auto": return "auto" if choice_type == "any": return "required" if choice_type == "tool": return {"type": "function", "function": {"name": choice.get("name", "")}} return "auto" def _parse_tool_arguments(arguments: Any) -> Any: """Parse tool call arguments from string to dict. LiteLLM/OpenAI returns arguments as a JSON string, but Anthropic expects input as a parsed dict. """ if isinstance(arguments, str): try: return json.loads(arguments) except (json.JSONDecodeError, TypeError): return arguments return arguments 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__( self, provider: str = "bedrock", region: str | None = None, profile_name: str | None = None, **kwargs: Any, ): """Initialize LiteLLM backend. Args: provider: LiteLLM provider prefix (bedrock, vertex_ai, openrouter, etc.) region: Cloud region (provider-specific) profile_name: AWS named profile for credential resolution (bedrock only). When set, boto3 uses this profile (e.g. an SSO profile) instead of the ambient credentials. Ignored for non-bedrock providers. **kwargs: Additional provider-specific config """ if not LITELLM_AVAILABLE: raise ImportError( "litellm is required for LiteLLMBackend. Install with: pip install litellm" ) self.provider = provider self.region = region self.profile_name = profile_name self.kwargs = kwargs # Get provider config from registry self._config = get_provider_config(provider) # For Bedrock, fetch model map dynamically from AWS API if provider == "bedrock": # litellm takes the botocore-backed `_auth_with_aws_session_token` # path as soon as temporary credentials (AWS_SESSION_TOKEN) are # present. botocore is an optional dependency (the `bedrock` # extra); when it is absent — as in the slim default Docker image — # the failure only surfaces at request time as a misleading # `authentication_error: No module named 'botocore'` (#1551). Fail # fast at startup with an actionable message instead. if os.environ.get("AWS_SESSION_TOKEN") and importlib.util.find_spec("botocore") is None: raise ImportError( "Bedrock with temporary credentials (AWS_SESSION_TOKEN) requires " "botocore, which is not installed. Install the bedrock extra: " "pip install 'headroom-ai[bedrock]' (or pip install botocore)." ) self._model_map = _fetch_bedrock_inference_profiles(region, profile_name=profile_name) litellm.set_verbose = False # Reduce noise else: self._model_map = self._config.model_map # Operator override map (all providers; only meaningful for Bedrock # today). Lets you pin a plain model name to a specific target the # AWS discovery can't disambiguate — e.g. a per-user application # inference profile ARN for cost attribution. See # `_parse_bedrock_model_overrides`. self._model_overrides = _parse_bedrock_model_overrides( os.environ.get("HEADROOM_BEDROCK_MODEL_MAP") ) if self._model_overrides: logger.info( f"Loaded {len(self._model_overrides)} Bedrock model override(s) " f"from HEADROOM_BEDROCK_MODEL_MAP: {sorted(self._model_overrides)}" ) logger.info(f"LiteLLM backend initialized (provider={provider}, region={region})") @property def name(self) -> str: return f"litellm-{self.provider}" def map_model_id(self, anthropic_model: str) -> str: """Map Anthropic model ID to LiteLLM model string. Handles various input formats: - "claude-sonnet-4-20250514" (standard Anthropic) - "anthropic.claude-sonnet-4-20250514-v1:0" (Bedrock without region) - "us.anthropic.claude-sonnet-4-20250514-v1:0" (Bedrock with region) - "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0" (LiteLLM format) - "arn:aws:bedrock:...:application-inference-profile/..." (application inference profile) """ # Operator override wins over everything — an explicit pin the AWS # discovery cannot express (e.g. a per-user application inference # profile). Keyed by the plain name Claude Code sends. override = self._model_overrides.get(anthropic_model) if override: if override.startswith("arn:aws:"): # Application inference profile ARNs must use the converse # route — the invoke route rejects ARNs with HTTP 400. return f"bedrock/converse/{override}" if override.startswith(f"{self.provider}/"): return override return f"{self.provider}/{override}" # Check direct mapping first if anthropic_model in self._model_map: return self._model_map[anthropic_model] # For Bedrock, try to normalize various input formats if self.provider == "bedrock": # Application inference profile ARNs must use the converse route — # the invoke route rejects ARNs with HTTP 400. if anthropic_model.startswith("arn:aws:"): return f"bedrock/converse/{anthropic_model}" normalized = _normalize_bedrock_profile_id(anthropic_model) if normalized and normalized in self._model_map: return self._model_map[normalized] # Bedrock fallback: construct a valid region-prefixed model ID. # Without this, bare model names like "claude-sonnet-4-20250514" # would become "bedrock/claude-sonnet-4-20250514" which is not a # valid Bedrock model identifier. if "/" not in anthropic_model and anthropic_model.startswith("claude"): region_prefix = _bedrock_region_prefix(self.region or "us-east-1") return f"bedrock/{region_prefix}.anthropic.{anthropic_model}-v1:0" # 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 # Fallback: construct provider/model format return f"{self.provider}/{anthropic_model}" 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]]: """Convert Anthropic message format to LiteLLM/OpenAI format. Anthropic and OpenAI have different representations for tool calls: - Anthropic: assistant content blocks with type=tool_use, user content blocks with type=tool_result - OpenAI: assistant message with tool_calls field, separate role=tool messages This method converts Anthropic-style messages to OpenAI-style so LiteLLM can send them to any provider. """ converted = [] for msg in messages: role = msg.get("role", "user") content = msg.get("content", "") # Handle string content directly if isinstance(content, str): converted.append({"role": role, "content": content}) continue # Handle content blocks (Anthropic style) if isinstance(content, list): # Separate blocks by type text_parts = [] tool_use_blocks = [] tool_result_blocks = [] for block in content: if not isinstance(block, dict): continue block_type = block.get("type", "") if block_type == "text": text_parts.append(block.get("text", "")) elif block_type == "tool_use": tool_use_blocks.append(block) elif block_type == "tool_result": tool_result_blocks.append(block) # tool_result blocks → OpenAI "tool" role messages if tool_result_blocks: # Do NOT insert a separate user text message here — Bedrock # requires tool role messages to appear immediately after the # assistant tool_calls message with no intervening messages. # Any text alongside tool_result is discarded (Claude Code # doesn't send text with tool_result blocks in practice). for tr in tool_result_blocks: tr_content = tr.get("content", "") if isinstance(tr_content, list): tr_content = "\n".join( b.get("text", "") for b in tr_content if b.get("type") == "text" ) converted.append( { "role": "tool", "tool_call_id": tr["tool_use_id"], "content": str(tr_content), } ) continue # tool_use blocks → OpenAI assistant message with tool_calls if tool_use_blocks: assistant_msg: dict[str, Any] = {"role": "assistant"} if text_parts: assistant_msg["content"] = "\n".join(text_parts) else: assistant_msg["content"] = None assistant_msg["tool_calls"] = [ { "id": tu["id"], "type": "function", "function": { "name": tu["name"], "arguments": json.dumps(tu.get("input", {})), }, } for tu in tool_use_blocks ] converted.append(assistant_msg) continue # Simple text only if text_parts: converted.append({"role": role, "content": "\n".join(text_parts)}) else: converted.append({"role": role, "content": ""}) return converted def _to_anthropic_response( self, litellm_response: Any, original_model: str, ) -> dict[str, Any]: """Convert LiteLLM/OpenAI response to Anthropic format.""" msg_id = f"msg_{uuid.uuid4().hex[:24]}" # Extract content from OpenAI format choice = litellm_response.choices[0] message = choice.message # Build Anthropic content blocks content = [] if message.content: content.append({"type": "text", "text": message.content}) # Handle tool calls if present if hasattr(message, "tool_calls") and message.tool_calls: for tc in message.tool_calls: content.append( { "type": "tool_use", "id": tc.id, "name": tc.function.name, "input": _parse_tool_arguments(tc.function.arguments), } ) # Map stop reason stop_reason_map = { "stop": "end_turn", "length": "max_tokens", "tool_calls": "tool_use", "content_filter": "end_turn", } stop_reason = stop_reason_map.get(choice.finish_reason, "end_turn") # Build usage usage = _anthropic_usage_from_litellm(litellm_response.usage) return { "id": msg_id, "type": "message", "role": "assistant", "content": content, "model": original_model, "stop_reason": stop_reason, "stop_sequence": None, "usage": usage, } async def send_message( self, body: dict[str, Any], headers: dict[str, str], ) -> BackendResponse: """Send message via LiteLLM.""" original_model = body.get("model", "claude-3-5-sonnet-20241022") litellm_model = self.map_model_id(original_model) try: # Convert messages messages = self._convert_messages_for_litellm(body.get("messages", [])) # Build kwargs for litellm kwargs: dict[str, Any] = { "model": litellm_model, "messages": messages, } # Optional parameters if "max_tokens" in body: kwargs["max_tokens"] = body["max_tokens"] if "temperature" in body: kwargs["temperature"] = body["temperature"] if "top_p" in body: kwargs["top_p"] = body["top_p"] if "stop_sequences" in body: kwargs["stop"] = body["stop_sequences"] # Tools (convert Anthropic format to OpenAI format) if "tools" in body: kwargs["tools"] = [_convert_anthropic_tool(t) for t in body["tools"]] if "tool_choice" in body: kwargs["tool_choice"] = _convert_tool_choice(body["tool_choice"]) # System prompt (Anthropic puts it in body, OpenAI in messages) if "system" in body: system = body["system"] if isinstance(system, str): kwargs["messages"].insert(0, {"role": "system", "content": system}) elif isinstance(system, list): # Anthropic list format system_text = " ".join( s.get("text", "") if isinstance(s, dict) else str(s) for s in system ) kwargs["messages"].insert(0, {"role": "system", "content": system_text}) # Provider-specific region config if self.region: if self.provider == "bedrock": kwargs["aws_region_name"] = self.region elif self.provider in ("vertex_ai", "vertex_ai_beta"): kwargs["vertex_location"] = self.region if self.provider == "bedrock" and self.profile_name: kwargs["aws_profile_name"] = self.profile_name # Forward API key from request headers if present. # Skip for Bedrock/Vertex: they use env-based auth (AWS SigV4 / Google ADC). # Forwarding x-api-key (e.g. sk-ant-dummy) would override their credentials. _env_auth_providers = ("bedrock", "vertex_ai", "vertex_ai_beta", "sagemaker") if self.provider not in _env_auth_providers: auth_header = headers.get("authorization", headers.get("Authorization", "")) if auth_header.startswith("Bearer "): kwargs["api_key"] = auth_header[7:] elif headers.get("x-api-key"): kwargs["api_key"] = headers["x-api-key"] logger.debug(f"LiteLLM request: model={litellm_model}") # Make the call response = await acompletion(**kwargs) # Convert to Anthropic format anthropic_response = self._to_anthropic_response(response, original_model) return BackendResponse( body=anthropic_response, status_code=200, headers={"content-type": "application/json"}, ) except Exception as e: logger.error(f"LiteLLM error: {e}") # Map to Anthropic 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 = "authentication_error" status_code = 401 elif "rate" in error_str or "limit" in error_str: error_type = "rate_limit_error" status_code = 429 elif "not found" in error_str: error_type = "not_found_error" status_code = 404 return BackendResponse( body={ "type": "error", "error": {"type": error_type, "message": str(e)}, }, status_code=status_code, error=str(e), ) async def stream_message( self, body: dict[str, Any], headers: dict[str, str], ) -> AsyncIterator[StreamEvent]: """Stream message via LiteLLM. Translates OpenAI streaming chunks into Anthropic SSE events. Handles both text content and tool_calls dynamically — block types are emitted based on what LiteLLM actually returns, not hardcoded. """ original_model = body.get("model", "claude-3-5-sonnet-20241022") litellm_model = self.map_model_id(original_model) try: messages = self._convert_messages_for_litellm(body.get("messages", [])) kwargs: dict[str, Any] = { "model": litellm_model, "messages": messages, "stream": True, } if "max_tokens" in body: kwargs["max_tokens"] = body["max_tokens"] if "temperature" in body: kwargs["temperature"] = body["temperature"] if "top_p" in body: kwargs["top_p"] = body["top_p"] if "stop_sequences" in body: kwargs["stop"] = body["stop_sequences"] if "tools" in body: kwargs["tools"] = [_convert_anthropic_tool(t) for t in body["tools"]] if "tool_choice" in body: kwargs["tool_choice"] = _convert_tool_choice(body["tool_choice"]) if "system" in body: system = body["system"] if isinstance(system, str): kwargs["messages"].insert(0, {"role": "system", "content": system}) elif isinstance(system, list): system_text = " ".join( s.get("text", "") if isinstance(s, dict) else str(s) for s in system ) kwargs["messages"].insert(0, {"role": "system", "content": system_text}) # Provider-specific region config if self.region: if self.provider == "bedrock": kwargs["aws_region_name"] = self.region elif self.provider in ("vertex_ai", "vertex_ai_beta"): kwargs["vertex_location"] = self.region if self.provider == "bedrock" and self.profile_name: kwargs["aws_profile_name"] = self.profile_name # Forward API key from request headers if present. # Skip for Bedrock/Vertex: they use env-based auth (AWS SigV4 / Google ADC). # Forwarding x-api-key (e.g. sk-ant-dummy) would override their credentials. _env_auth_providers = ("bedrock", "vertex_ai", "vertex_ai_beta", "sagemaker") if self.provider not in _env_auth_providers: auth_header = headers.get("authorization", headers.get("Authorization", "")) if auth_header.startswith("Bearer "): kwargs["api_key"] = auth_header[7:] elif headers.get("x-api-key"): kwargs["api_key"] = headers["x-api-key"] msg_id = f"msg_{uuid.uuid4().hex[:24]}" # Emit message_start yield StreamEvent( event_type="message_start", data={ "type": "message_start", "message": { "id": msg_id, "type": "message", "role": "assistant", "content": [], "model": original_model, "stop_reason": None, "stop_sequence": None, "usage": {"input_tokens": 0, "output_tokens": 0}, }, }, ) # Stream content — blocks emitted dynamically based on response response = await acompletion(**kwargs) output_tokens = 0 current_block_index = -1 active_block_type: str | None = None # "text" or "tool_use" tool_block_map: dict[int, int] = {} # litellm tc.index → SSE block index stop_reason = "end_turn" async for chunk in response: if not hasattr(chunk, "choices") or not chunk.choices: continue choice = chunk.choices[0] delta = choice.delta # Check finish_reason to set stop_reason if choice.finish_reason == "tool_calls": stop_reason = "tool_use" elif choice.finish_reason == "stop": stop_reason = "end_turn" elif choice.finish_reason == "length": stop_reason = "max_tokens" # Handle tool_calls in the delta if hasattr(delta, "tool_calls") and delta.tool_calls: for tc in delta.tool_calls: idx = tc.index if tc.index is not None else 0 if idx not in tool_block_map: # Close previous block if open if active_block_type is not None: yield StreamEvent( event_type="content_block_stop", data={ "type": "content_block_stop", "index": current_block_index, }, ) # Open a new tool_use block current_block_index += 1 tool_block_map[idx] = current_block_index active_block_type = "tool_use" tool_id = tc.id or f"toolu_{uuid.uuid4().hex[:24]}" tool_name = tc.function.name if tc.function and tc.function.name else "" yield StreamEvent( event_type="content_block_start", data={ "type": "content_block_start", "index": current_block_index, "content_block": { "type": "tool_use", "id": tool_id, "name": tool_name, "input": {}, }, }, ) # Emit argument deltas if tc.function and tc.function.arguments: block_idx = tool_block_map[idx] yield StreamEvent( event_type="content_block_delta", data={ "type": "content_block_delta", "index": block_idx, "delta": { "type": "input_json_delta", "partial_json": tc.function.arguments, }, }, ) output_tokens += 1 # Handle text content in the delta elif hasattr(delta, "content") and delta.content: if active_block_type != "text": # Close previous block if open if active_block_type is not None: yield StreamEvent( event_type="content_block_stop", data={ "type": "content_block_stop", "index": current_block_index, }, ) # Open a new text block current_block_index += 1 active_block_type = "text" yield StreamEvent( event_type="content_block_start", data={ "type": "content_block_start", "index": current_block_index, "content_block": {"type": "text", "text": ""}, }, ) yield StreamEvent( event_type="content_block_delta", data={ "type": "content_block_delta", "index": current_block_index, "delta": {"type": "text_delta", "text": delta.content}, }, ) output_tokens += 1 # Close the last open block if active_block_type is not None: yield StreamEvent( event_type="content_block_stop", data={"type": "content_block_stop", "index": current_block_index}, ) # Emit message_delta with correct stop reason yield StreamEvent( event_type="message_delta", data={ "type": "message_delta", "delta": {"stop_reason": stop_reason, "stop_sequence": None}, "usage": {"output_tokens": output_tokens}, }, ) # Emit message_stop yield StreamEvent( event_type="message_stop", data={"type": "message_stop"}, ) except Exception as e: logger.error(f"LiteLLM streaming error: {e}") yield StreamEvent( event_type="error", data={ "type": "error", "error": {"type": "api_error", "message": str(e)}, }, ) 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 region config if self.region: if self.provider == "bedrock": kwargs["aws_region_name"] = self.region elif self.provider in ("vertex_ai", "vertex_ai_beta"): kwargs["vertex_location"] = self.region if self.provider == "bedrock" and self.profile_name: kwargs["aws_profile_name"] = self.profile_name # Forward API key from request headers if present. # Skip for Bedrock/Vertex: they use env-based auth (AWS SigV4 / Google ADC). # Forwarding x-api-key (e.g. sk-ant-dummy) would override their credentials. _env_auth_providers = ("bedrock", "vertex_ai", "vertex_ai_beta", "sagemaker") if self.provider not in _env_auth_providers: auth_header = headers.get("authorization", headers.get("Authorization", "")) if auth_header.startswith("Bearer "): kwargs["api_key"] = auth_header[7:] elif headers.get("x-api-key"): kwargs["api_key"] = headers["x-api-key"] logger.debug(f"LiteLLM OpenAI request: model={litellm_model}") # Make the call response = await acompletion(**kwargs) # Build the usage block. LiteLLM normalizes prompt-cache stats from # multiple providers (Anthropic, Bedrock-Claude, OpenAI prompt-caching, # DeepSeek) onto its Usage object — top-level # cache_read_input_tokens / cache_creation_input_tokens for the # Anthropic-style dialect, and prompt_tokens_details.cached_tokens / # cache_creation_tokens for the OpenAI nested dialect. Surface both # so PrefixCacheTracker.update_from_response on the backend-routed # path observes a stable shape instead of branching on key presence. usage_block: dict[str, Any] = { "prompt_tokens": response.usage.prompt_tokens, "completion_tokens": response.usage.completion_tokens, "total_tokens": response.usage.total_tokens, } # Defensive getattr: LiteLLM only attaches these top-level attrs # when the underlying provider returned cache stats. Zero is the # cold-start / no-cache value. cache_read = int(getattr(response.usage, "cache_read_input_tokens", 0) or 0) cache_write = int(getattr(response.usage, "cache_creation_input_tokens", 0) or 0) # OpenAI nested dialect — fall back here if the top-level dialect is # absent (pure OpenAI prompt caching). ptd_obj = getattr(response.usage, "prompt_tokens_details", None) ptd_cached = 0 ptd_cache_creation = 0 if ptd_obj is not None: ptd_cached = int(getattr(ptd_obj, "cached_tokens", 0) or 0) ptd_cache_creation = int(getattr(ptd_obj, "cache_creation_tokens", 0) or 0) final_cache_read = cache_read or ptd_cached final_cache_write = cache_write or ptd_cache_creation if final_cache_read or final_cache_write: usage_block["cache_read_input_tokens"] = final_cache_read usage_block["cache_creation_input_tokens"] = final_cache_write # Mirror into the OpenAI nested shape so callers that only know # the OpenAI dialect can read it without branching. usage_block["prompt_tokens_details"] = {"cached_tokens": final_cache_read} logger.debug( f"LiteLLM OpenAI cache stats: cache_read={final_cache_read} " f"cache_write={final_cache_write} model={litellm_model}" ) # 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": usage_block, } 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), ) async def stream_openai_message( self, body: dict[str, Any], headers: dict[str, str], ) -> AsyncIterator[str]: """Stream OpenAI-format chat completion via LiteLLM. Yields SSE-formatted strings ready to send to the client. """ original_model = body.get("model", "gpt-4") litellm_model = self.map_model_id(original_model) try: kwargs: dict[str, Any] = { "model": litellm_model, "messages": body.get("messages", []), "stream": True, } for param in [ "max_tokens", "temperature", "top_p", "stop", "tools", "tool_choice", "response_format", "seed", "n", ]: if param in body: kwargs[param] = body[param] if "stream_options" in body: kwargs["stream_options"] = body["stream_options"] # Provider-specific region config if self.region: if self.provider == "bedrock": kwargs["aws_region_name"] = self.region elif self.provider in ("vertex_ai", "vertex_ai_beta"): kwargs["vertex_location"] = self.region if self.provider == "bedrock" and self.profile_name: kwargs["aws_profile_name"] = self.profile_name # Forward API key from request headers if present. # Skip for Bedrock/Vertex: they use env-based auth (AWS SigV4 / Google ADC). # Forwarding x-api-key (e.g. sk-ant-dummy) would override their credentials. _env_auth_providers = ("bedrock", "vertex_ai", "vertex_ai_beta", "sagemaker") if self.provider not in _env_auth_providers: auth_header = headers.get("authorization", headers.get("Authorization", "")) if auth_header.startswith("Bearer "): kwargs["api_key"] = auth_header[7:] elif headers.get("x-api-key"): kwargs["api_key"] = headers["x-api-key"] response = await acompletion(**kwargs) async for chunk in response: chunk_dict = chunk.model_dump(exclude_none=True, exclude_unset=True) yield f"data: {json.dumps(chunk_dict)}\n\n" yield "data: [DONE]\n\n" except Exception as e: logger.error(f"LiteLLM OpenAI streaming error: {e}") error_data = { "error": { "message": str(e), "type": "api_error", "code": "backend_error", } } yield f"data: {json.dumps(error_data)}\n\n" yield "data: [DONE]\n\n"