Merge remote-tracking branch 'origin/main' into litellm_patch_model_reject_unloadable_params

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	litellm/router.py
This commit is contained in:
Moe Khalil 2026-10-26 23:06:00 +00:00
commit 447ae4d472
106 changed files with 6013 additions and 2817 deletions

View file

@ -1,319 +0,0 @@
import asyncio
import aiohttp
import json
import math
from typing import Any
# Asynchronously fetch data from a given URL
async def fetch_data(url):
try:
# Create an asynchronous session
async with aiohttp.ClientSession() as session:
# Send a GET request to the URL
async with session.get(url) as resp:
# Raise an error if the response status is not OK
resp.raise_for_status()
# Parse the response JSON
resp_json = await resp.json()
print("Fetch the data from URL.")
# Return the 'data' field from the JSON response
return resp_json['data']
except Exception as e:
# Print an error message if fetching data fails
print("Error fetching data from URL:", e)
return None
FRIENDLI_API_URL = "https://api.friendli.ai/serverless/v1/models"
FRIENDLI_PROVIDER = "friendliai"
INHERITABLE_BASE_KEYS = (
"supports_pdf_input",
"supports_assistant_prefill",
"supports_adaptive_thinking",
"supports_output_config",
)
REASONING_EFFORT_LEVEL_ORDER = ("none", "minimal", "low", "medium", "high", "xhigh", "max")
def _find_base_model_entry(base_model: str, local_data: dict) -> str | None:
if not base_model:
return None
bm_tail = base_model.split("/")[-1].lower()
if base_model in local_data:
return base_model
for key in local_data:
if key.startswith("sample_spec") or key == "fallback_generalizations":
continue
if key.split("/")[-1].lower() == bm_tail:
return key
return None
def _reasoning_effort_levels(reasoning_options: list) -> list:
offered = {
val
for opt in reasoning_options or []
if opt.get("type") == "effort"
for val in opt.get("values", [])
}
return [level for level in REASONING_EFFORT_LEVEL_ORDER if level in offered]
def _valid_token_price(value: object) -> bool:
try:
price = float(value) # pyright: ignore[reportArgumentType] # non-numeric values are rejected via the except
except (TypeError, ValueError):
return False
return math.isfinite(price) and price >= 0
def _has_valid_token_prices(pricing: dict | None) -> bool:
prices = pricing or {}
return _valid_token_price(prices.get("input")) and _valid_token_price(prices.get("output"))
def _pricing(pricing: dict) -> dict:
out: dict[str, Any] = {}
if not pricing:
return out
if "input" in pricing:
out["input_cost_per_token"] = float(pricing["input"])
if "output" in pricing:
out["output_cost_per_token"] = float(pricing["output"])
if "input_cache_read" in pricing and pricing["input_cache_read"] is not None:
out["cache_read_input_token_cost"] = float(pricing["input_cache_read"])
return out
def _modality_flags(input_mods: list) -> dict:
mods = input_mods or []
has_image = "image" in mods
return {
"supports_vision": has_image,
"supports_image_input": has_image,
"supports_video_input": "video" in mods,
}
def transform_friendli_data(data: list, local_data: dict) -> dict:
transformed: dict[str, dict] = {}
if not data:
return transformed
for model in data:
# An unpriced row must never wholesale-replace an already priced local entry:
# missing prices cost-calculate as zero, silently zeroing tracked spend
if not _has_valid_token_prices(model.get("pricing")):
continue
model_id = model["id"]
base_model = model.get("base_model") or ""
entry: dict[str, Any] = {
"litellm_provider": FRIENDLI_PROVIDER,
}
base_key = _find_base_model_entry(base_model, local_data)
if base_key:
base_entry = local_data[base_key]
for k in INHERITABLE_BASE_KEYS:
if k in base_entry:
entry[k] = base_entry[k]
ctx = model.get("context_length")
if ctx is not None:
entry["max_input_tokens"] = int(ctx)
max_out = model.get("max_completion_tokens")
if max_out is not None:
entry["max_output_tokens"] = int(max_out)
entry["max_tokens"] = int(max_out)
pricing = _pricing(model.get("pricing", {}))
entry.update(pricing)
entry["supports_prompt_caching"] = "cache_read_input_token_cost" in pricing
reasoning = model.get("reasoning") is True
entry["supports_reasoning"] = reasoning
if reasoning:
entry["reasoning_effort_levels"] = _reasoning_effort_levels(
model.get("reasoning_options", [])
)
func = model.get("functionality", {})
entry["supports_function_calling"] = func.get("tool_call") is True
entry["supports_parallel_function_calling"] = func.get("parallel_tool_call") is True
is_struct = func.get("structured_output") is True
entry["supports_response_schema"] = is_struct
entry["supports_native_structured_output"] = is_struct
entry["supports_system_messages"] = func.get("system_messages") is True
entry["supports_tool_choice"] = func.get("tool_choice") is True
entry.update(_modality_flags(model.get("input_modalities", [])))
entry["mode"] = model.get("mode", "chat")
desc = model.get("description")
if desc:
entry["comment"] = desc
dep = model.get("deprecation_date")
if dep:
entry["deprecation_date"] = dep.split("T")[0]
entry["source"] = FRIENDLI_API_URL
transformed[f"{FRIENDLI_PROVIDER}/{model_id}"] = entry
return transformed
# Synchronize local data with remote data
def sync_local_data_with_remote(local_data, remote_data, replace_keys=frozenset()):
# Update existing keys in local_data with values from remote_data
# (replace_keys entries are swapped wholesale so a field the remote catalog
# dropped, e.g. cache pricing, cannot survive as a stale value)
for key in (set(local_data) & set(remote_data)):
if key in replace_keys:
local_data[key] = remote_data[key]
else:
local_data[key].update(remote_data[key])
# Add new keys from remote_data to local_data
for key in (set(remote_data) - set(local_data)):
local_data[key] = remote_data[key]
# Write data to the json file
def write_to_file(file_path, data):
try:
# Open the file in write mode
with open(file_path, "w") as file:
# Dump the data as JSON into the file
json.dump(data, file, indent=4)
print("Values updated successfully.")
except Exception as e:
# Print an error message if writing to file fails
print("Error updating JSON file:", e)
# Update the existing models and add the missing models for OpenRouter
def transform_openrouter_data(data):
transformed = {}
if not data:
return transformed
for row in data:
# Add the fields 'max_tokens' and 'input_cost_per_token'
obj = {
"max_tokens": row["context_length"],
"input_cost_per_token": float(row["pricing"]["prompt"]),
}
# Add 'max_output_tokens' as a field if it is not None
if "top_provider" in row and "max_completion_tokens" in row["top_provider"] and row["top_provider"]["max_completion_tokens"] is not None:
obj['max_output_tokens'] = int(row["top_provider"]["max_completion_tokens"])
# Add the field 'output_cost_per_token'
obj.update({
"output_cost_per_token": float(row["pricing"]["completion"]),
})
# Add field 'input_cost_per_image' if it exists and is non-zero
if "pricing" in row and "image" in row["pricing"] and float(row["pricing"]["image"]) != 0.0:
obj['input_cost_per_image'] = float(row["pricing"]["image"])
# Add the fields 'litellm_provider' and 'mode'
obj.update({
"litellm_provider": "openrouter",
"mode": "chat"
})
# Add the 'supports_vision' field if the modality is 'multimodal'
if row.get('architecture', {}).get('modality') == 'multimodal':
obj['supports_vision'] = True
# Use a composite key to store the transformed object
transformed[f'openrouter/{row["id"]}'] = obj
return transformed
# Update the existing models and add the missing models for Vercel AI Gateway
def transform_vercel_ai_gateway_data(data):
transformed = {}
if not data:
return transformed
for row in data:
# Rows without token pricing or token limits (video/embedding models) previously KeyError'd the whole sync
if any(row.get(k) is None for k in ("context_window", "max_tokens")) or any(
row.get("pricing", {}).get(k) is None for k in ("input", "output")
):
continue
obj = {
"max_tokens": row["context_window"],
"input_cost_per_token": float(row["pricing"]["input"]),
"output_cost_per_token": float(row["pricing"]["output"]),
'max_output_tokens': row['max_tokens'],
'max_input_tokens': row["context_window"],
}
# Handle cache pricing if available
if "pricing" in row:
if "input_cache_read" in row["pricing"] and row["pricing"]["input_cache_read"] is not None:
obj['cache_read_input_token_cost'] = float(f"{float(row['pricing']['input_cache_read']):e}")
if "input_cache_write" in row["pricing"] and row["pricing"]["input_cache_write"] is not None:
obj['cache_creation_input_token_cost'] = float(f"{float(row['pricing']['input_cache_write']):e}")
mode = "embedding" if "embedding" in row["id"].lower() else "chat"
obj.update({"litellm_provider": "vercel_ai_gateway", "mode": mode})
transformed[f'vercel_ai_gateway/{row["id"]}'] = obj
return transformed
# Load local data from a specified file
def load_local_data(file_path):
try:
# Open the file in read mode
with open(file_path, "r") as file:
# Load and return the JSON data
return json.load(file)
except FileNotFoundError:
# Print an error message if the file is not found
print("File not found:", file_path)
return None
except json.JSONDecodeError as e:
# Print an error message if JSON decoding fails
print("Error decoding JSON:", e)
return None
def main():
local_file_path = "model_prices_and_context_window.json" # Path to the local data file
openrouter_url = "https://openrouter.ai/api/v1/models" # URL to fetch OpenRouter data
vercel_ai_gateway_url = "https://ai-gateway.vercel.sh/v1/models" # URL to fetch Vercel AI Gateway data
# Load local data from file
local_data = load_local_data(local_file_path)
# Fetch OpenRouter data
openrouter_data = asyncio.run(fetch_data(openrouter_url))
# Transform the fetched OpenRouter data
openrouter_data = transform_openrouter_data(openrouter_data)
# Fetch Vercel AI Gateway data
vercel_data = asyncio.run(fetch_data(vercel_ai_gateway_url))
# Transform the fetched Vercel AI Gateway data
vercel_data = transform_vercel_ai_gateway_data(vercel_data)
friendli_data = asyncio.run(fetch_data(FRIENDLI_API_URL))
friendli_data = transform_friendli_data(friendli_data, local_data)
# Combine both datasets
all_remote_data = {**openrouter_data, **vercel_data, **friendli_data}
# If both local and openrouter data are available, synchronize and save
if local_data and all_remote_data:
sync_local_data_with_remote(local_data, all_remote_data, replace_keys=frozenset(friendli_data))
write_to_file(local_file_path, local_data)
else:
print("Failed to fetch model data from either local file or URL.")
# Entry point of the script
if __name__ == "__main__":
main()

View file

@ -1,36 +0,0 @@
name: Updates model_prices_and_context_window.json and Create Pull Request
on:
schedule:
- cron: "0 0 * * 0" # Run every Sundays at midnight
#- cron: "0 0 * * *" # Run daily at midnight
permissions:
contents: write
pull-requests: write
jobs:
auto_update_price_and_context_window:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Update JSON Data
run: |
uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py"
- name: Create Pull Request
run: |
git add model_prices_and_context_window.json
git commit -m "Update model_prices_and_context_window.json file: $(date +'%Y-%m-%d')"
gh pr create --title "Update model_prices_and_context_window.json file" \
--body "Automated update for model_prices_and_context_window.json" \
--head auto-update-price-and-context-window-$(date +'%Y-%m-%d') \
--base main
env:
GH_TOKEN: ${{ secrets.GH_TOKEN }}

View file

@ -1,65 +0,0 @@
name: Sync Together AI model registry
on:
schedule:
- cron: "30 6 * * *"
workflow_dispatch:
permissions:
contents: write
pull-requests: write
jobs:
sync_together_ai_models:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
env:
BASE_BRANCH: ${{ github.event.repository.default_branch }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: ${{ env.BASE_BRANCH }}
persist-credentials: false
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Look for an already-open sync PR
id: existing
run: |
open_pr="$(gh pr list --repo "$GITHUB_REPOSITORY" --state open --limit 1000 --json headRefName \
--jq '[.[].headRefName | select(startswith("litellm_together_registry_sync_"))] | first // empty')"
echo "open_pr=$open_pr" >> "$GITHUB_OUTPUT"
if [ -n "$open_pr" ]; then
echo "An open sync PR already exists on branch $open_pr; skipping this run."
fi
env:
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}
- name: Run the sync
if: steps.existing.outputs.open_pr == ''
run: |
uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md"
env:
TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }}
- name: Create a pull request when the registry changed
if: steps.existing.outputs.open_pr == ''
run: |
if git diff --quiet; then
echo "Registry already in sync; no PR needed."
exit 0
fi
branch="litellm_together_registry_sync_$(date +'%Y-%m-%d')"
git config user.name "github-actions[bot]"
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
git checkout -b "$branch"
git add model_prices_and_context_window.json \
litellm/model_prices_and_context_window_backup.json
git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')"
gh auth setup-git
git push origin "$branch"
gh pr create --title "feat(models): sync together_ai model registry" \
--body-file "$RUNNER_TEMP/pr_body.md" \
--head "$branch" \
--base "$BASE_BRANCH"
env:
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}

View file

@ -0,0 +1,76 @@
# Use your ChatGPT subscription through LiteLLM
Sign into ChatGPT in Codex or Pi and send model requests through your existing LiteLLM proxy. The client keeps your login and refreshes it. Each request uses that client's ChatGPT account, while a separate LiteLLM virtual key controls gateway access and usage attribution.
An administrator adds a `chatgpt/` model using the usual model configuration. No ChatGPT login needs to be added to the proxy. This uses the same separate-key convention as provider passthrough, including Claude Code subscription forwarding.
## Add a model to the proxy
Reuse your existing proxy. If you need one, follow the [proxy quickstart](https://docs.litellm.ai/docs/proxy/docker_quick_start). Merge this entry into its model list and restart the proxy:
```yaml
model_list:
- model_name: gpt-6-luna
litellm_params:
model: chatgpt/gpt-6-luna
model_info:
mode: responses
```
Use a model available to the caller's subscription. The virtual key must allow the exposed model name. Leave saved ChatGPT logins and API credentials unset for a client-only deployment. Explicit cross-provider fallbacks can incur separate API charges; do not configure them for subscription-only traffic.
## Configure Codex
Use your existing LiteLLM URL and a virtual key assigned to your gateway user. Export that key as `LITELLM_API_KEY` in the environment that launches Codex. Keep the value out of committed configuration.
Add this provider to `~/.codex/config.toml`, preserving your other settings:
```toml
model = "gpt-6-luna"
model_provider = "litellm_subscription"
[model_providers.litellm_subscription]
name = "LiteLLM subscription"
base_url = "https://litellm.example.com/v1"
wire_api = "responses"
supports_websockets = false
requires_openai_auth = true
env_http_headers = { "x-litellm-api-key" = "LITELLM_API_KEY" }
```
Replace `https://litellm.example.com` with your proxy URL. A local proxy normally uses `http://localhost:4000/v1`. Launch Codex, sign into ChatGPT if needed, and send a prompt with a model available to your subscription. The selected model must match the name exposed by the gateway.
For a managed workspace rollout, use OpenAI's [gateway configuration guide](https://learn.chatgpt.com/docs/enterprise/sign-in-with-chatgpt-through-a-gateway).
## Native compatibility and Pi
For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials but does not need a model deployment or resolve model aliases. Restricted keys must allow the upstream model name.
### Configure Pi
Sign into the built-in `openai-codex` provider in Pi. Merge this override into `~/.pi/agent/models.json` and export your LiteLLM virtual key as `LITELLM_API_KEY` before launching Pi:
```json
{
"providers": {
"openai-codex": {
"baseUrl": "https://litellm.example.com/backend-api/codex",
"headers": {
"x-litellm-api-key": "LITELLM_API_KEY"
}
}
}
}
```
Keep Pi's subscription login as the provider credential. The header value names an environment variable, not a key to paste into the file.
## Verify account usage
Send a prompt and check the completed request in LiteLLM's Logs page. Its gateway user/key and token counts should match the caller. Give each person their own LiteLLM virtual key for separate attribution. The gateway credential and ChatGPT login are independently authenticated; LiteLLM does not establish that both identities belong to the same person.
The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key must allow that name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions.
The proxy forwards the ChatGPT bearer and account header to the fixed OpenAI Codex endpoint. It removes gateway credentials and cookies from the upstream request and redacts credential headers in logging callbacks. The client remains responsible for refresh and reauthentication. Native passthrough returns authentication and quota errors without switching subscriptions. Standard model routes retain the caller login across ChatGPT retries; explicitly configured cross-provider fallbacks use their own credentials. An explicit workspace header is preserved for OpenAI to validate.
This flow does not connect accounts through the dashboard or pool subscription allowances. Use a Codex build that supports ChatGPT authentication with a custom model provider. The integration forwards HTTP and SSE requests; it does not add native WebSocket or other Codex backend routes.

View file

@ -76,6 +76,7 @@ OPTIONAL_KWARGS_KEYS: Final = (
"itpm",
"otpm",
"use_xai_oauth",
"chatgpt_auth_file",
"github_copilot_auth_type",
"github_copilot_host",
"github_copilot_user_session",

View file

@ -158,6 +158,14 @@ def declared_authenticating_provider(model: str | None, custom_llm_provider: str
return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None
def resolve_model_and_provider_for_metadata(model: str, custom_llm_provider: str | None = None) -> tuple[str, str]:
declared: Final = declared_authenticating_provider(model, custom_llm_provider)
if declared is not None:
return model.removeprefix(f"{declared}/"), declared
resolved_model, resolved_provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider)
return resolved_model, resolved_provider
def inferred_provider(model: str | None) -> str | None:
if not model:
return None
@ -842,7 +850,9 @@ def _get_openai_compatible_provider_info(
api_base,
dynamic_api_key,
custom_llm_provider,
) = litellm.ChatGPTConfig().get_openai_compatible_provider_info(model, api_base, api_key, custom_llm_provider)
) = litellm.ChatGPTConfig().get_openai_compatible_provider_info(
model, api_base, api_key, custom_llm_provider, litellm_params
)
elif custom_llm_provider == "novita":
api_base = api_base or get_secret("NOVITA_API_BASE") or "https://api.novita.ai/v3/openai"
dynamic_api_key = api_key or get_secret_str("NOVITA_API_KEY")

View file

@ -1,10 +1,23 @@
from collections.abc import Mapping, Sequence
from typing import Final
from pydantic import TypeAdapter
from litellm.types.utils import ProviderSpecificHeader
_PROVIDER_HEADERS: Final[TypeAdapter[ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None]] = TypeAdapter(
ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None
)
class ProviderSpecificHeaderUtils:
@staticmethod
def forwarded_authorization(provider_specific_header: object, custom_llm_provider: str | None) -> str | None:
headers: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers(
_PROVIDER_HEADERS.validate_python(provider_specific_header), custom_llm_provider
)
return next((value for key, value in headers.items() if key.lower() == "authorization"), None)
@staticmethod
def get_provider_specific_headers(
provider_specific_header: ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None,

View file

@ -1490,6 +1490,7 @@ def convert_to_gemini_tool_call_result(
message: ChatCompletionToolMessage | ChatCompletionFunctionMessage,
last_message_with_tool_calls: dict | None,
forward_function_call_id: bool = False,
supports_multimodal_function_responses: bool = True,
) -> VertexPartType | list[VertexPartType]:
"""
OpenAI message with a tool result looks like:
@ -1652,9 +1653,9 @@ def convert_to_gemini_tool_call_result(
_part: Final[VertexPartType] = {"function_response": _function_response}
# For multimodal function responses, Gemini expects media parts nested
# inside functionResponse.parts instead of sibling content parts.
if inline_data_list:
if not supports_multimodal_function_responses:
return [_part, *(VertexPartType(inline_data=inline_data) for inline_data in inline_data_list)]
_function_response["parts"] = [{"inline_data": inline_data} for inline_data in inline_data_list]
return [_part]

View file

@ -1,9 +1,9 @@
import base64
import json
import os
import time
from collections.abc import Mapping
from typing import Final, TypeAlias
from functools import lru_cache
from typing import TYPE_CHECKING, Final
import httpx
from pydantic import JsonValue, TypeAdapter, ValidationError
@ -24,19 +24,21 @@ from .common_utils import (
CHATGPT_OAUTH_TOKEN_URL,
GetAccessTokenError,
GetDeviceCodeError,
JsonObject,
RefreshAccessTokenError,
decode_jwt_claims,
extract_chatgpt_account_id,
)
if TYPE_CHECKING:
from litellm.types.router import GenericLiteLLMParams
TOKEN_EXPIRY_SKEW_SECONDS: Final = 60
TOKEN_REFRESH_TIMEOUT_SECONDS: Final = 30
DEVICE_CODE_TIMEOUT_SECONDS: Final = 15 * 60
DEVICE_CODE_COOLDOWN_SECONDS: Final = 5 * 60
DEVICE_CODE_POLL_SLEEP_SECONDS: Final = 5
OPENAI_AUTH_CLAIM_KEY: Final = "https://api.openai.com/auth"
JsonObject: TypeAlias = Mapping[str, JsonValue]
_JSON_OBJECT_ADAPTER: Final = TypeAdapter(JsonObject)
@ -44,6 +46,26 @@ def _optional_str(value: JsonValue | None) -> str | None:
return value if isinstance(value, str) else None
def get_chatgpt_auth_file(
litellm_params: "Mapping[str, object] | GenericLiteLLMParams | None",
) -> str | None:
if litellm_params is None:
return None
value: Final = (
litellm_params.get("chatgpt_auth_file")
if isinstance(litellm_params, Mapping)
else litellm_params.chatgpt_auth_file
)
if not isinstance(value, str) or not value:
return None
return os.path.abspath(os.path.expanduser(value))
@lru_cache(maxsize=128)
def get_cached_authenticator(auth_file: str) -> "Authenticator":
return Authenticator(auth_file=auth_file)
def _token_refresh_timeout() -> httpx.Timeout:
configured: Final = get_configured_request_timeout()
seconds: Final = (
@ -53,13 +75,13 @@ def _token_refresh_timeout() -> httpx.Timeout:
class Authenticator:
def __init__(self) -> None:
self.token_dir = os.getenv(
"CHATGPT_TOKEN_DIR",
os.path.expanduser("~/.config/litellm/chatgpt"),
def __init__(self, auth_file: str | None = None) -> None:
default_auth_file: Final = os.path.join(
os.getenv("CHATGPT_TOKEN_DIR", os.path.expanduser("~/.config/litellm/chatgpt")),
os.getenv("CHATGPT_AUTH_FILE", "auth.json"),
)
self.auth_file = os.path.join(self.token_dir, os.getenv("CHATGPT_AUTH_FILE", "auth.json"))
self._ensure_token_dir()
self.auth_file = os.path.abspath(os.path.expanduser(auth_file or default_auth_file))
self.token_dir = os.path.dirname(self.auth_file)
def get_api_base(self) -> str:
return os.getenv("CHATGPT_API_BASE") or os.getenv("OPENAI_CHATGPT_API_BASE") or CHATGPT_API_BASE
@ -99,16 +121,18 @@ class Authenticator:
tokens: Final = self._login_device_code()
return tokens["access_token"]
def get_account_id(self) -> str | None:
def get_account_id(self, access_token: str | None = None) -> str | None:
auth_data: Final = self._read_auth_file()
if access_token and (not auth_data or auth_data.get("access_token") != access_token):
return extract_chatgpt_account_id(access_token)
if not auth_data:
return None
account_id: Final = _optional_str(auth_data.get("account_id"))
if account_id:
return account_id
id_token: Final = auth_data.get("id_token")
access_token: Final = auth_data.get("access_token")
derived: Final = self._extract_account_id(_optional_str(id_token or access_token))
saved_access_token: Final = auth_data.get("access_token")
derived: Final = extract_chatgpt_account_id(_optional_str(id_token or saved_access_token))
if derived:
self._write_auth_file({**auth_data, "account_id": derived})
return derived
@ -129,6 +153,7 @@ class Authenticator:
def _write_auth_file(self, data: JsonObject) -> None:
try:
self._ensure_token_dir()
with open(self.auth_file, "w") as f:
json.dump(data, f)
except OSError as exc:
@ -146,34 +171,11 @@ class Authenticator:
return time.time() >= float(derived_expires_at) - TOKEN_EXPIRY_SKEW_SECONDS
def _get_expires_at(self, token: str) -> int | None:
claims: Final = self._decode_jwt_claims(token)
exp: Final = claims.get("exp")
exp: Final = decode_jwt_claims(token).get("exp")
if isinstance(exp, (int, float)):
return int(exp)
return None
def _decode_jwt_claims(self, token: str) -> JsonObject:
try:
parts: Final = token.split(".")
if len(parts) < 2:
return {}
payload_b64: Final = parts[1] + "=" * (-len(parts[1]) % 4)
payload_bytes: Final = base64.urlsafe_b64decode(payload_b64)
return _JSON_OBJECT_ADAPTER.validate_python(json.loads(payload_bytes.decode("utf-8")))
except Exception:
return {}
def _extract_account_id(self, token: str | None) -> str | None:
if not token:
return None
claims: Final = self._decode_jwt_claims(token)
auth_claims: Final = claims.get(OPENAI_AUTH_CLAIM_KEY)
if isinstance(auth_claims, dict):
account_id: Final = auth_claims.get("chatgpt_account_id")
if isinstance(account_id, str) and account_id:
return account_id
return None
def _login_device_code(self) -> dict[str, str]:
cooldown_remaining: Final = self._get_device_code_cooldown_remaining(self._read_auth_file())
if cooldown_remaining > 0:
@ -361,15 +363,16 @@ class Authenticator:
"refresh_token": _optional_str(data.get("refresh_token")) or refresh_token,
"id_token": id_token,
}
account_id: Final = self.get_account_id()
auth_data: Final = self._build_auth_record(refreshed)
self._write_auth_file(auth_data)
self._write_auth_file({**auth_data, "account_id": account_id} if account_id else auth_data)
return refreshed
def _build_auth_record(self, tokens: dict[str, str]) -> JsonObject:
access_token: Final = tokens.get("access_token")
id_token: Final = tokens.get("id_token")
expires_at: Final = self._get_expires_at(access_token) if access_token else None
account_id: Final = self._extract_account_id(id_token or access_token)
account_id: Final = extract_chatgpt_account_id(id_token or access_token)
return {
"access_token": access_token,
"refresh_token": tokens.get("refresh_token"),

View file

@ -1,14 +1,20 @@
from collections.abc import Mapping
from typing import Any, Final
from litellm.exceptions import AuthenticationError
from litellm.llms.openai.openai import OpenAIConfig
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
from ..authenticator import Authenticator
from ..authenticator import Authenticator, get_cached_authenticator, get_chatgpt_auth_file
from ..common_utils import (
GetAccessTokenError,
chatgpt_client_api_base,
ensure_chatgpt_session_id,
get_chatgpt_client_credential,
get_chatgpt_default_headers,
get_chatgpt_request_credential,
strip_chatgpt_client_credential_headers,
)
from .streaming_utils import ChatGPTToolCallNormalizer
@ -19,9 +25,27 @@ class ChatGPTConfig(OpenAIConfig):
api_key: str | None = None,
api_base: str | None = None,
custom_llm_provider: str = "openai",
authenticator: Authenticator | None = None,
) -> None:
super().__init__()
self.authenticator = Authenticator()
self.authenticator = authenticator if authenticator is not None else Authenticator()
def _resolve_authenticator(
self, litellm_params: Mapping[str, object] | GenericLiteLLMParams | None
) -> Authenticator:
auth_file: Final = get_chatgpt_auth_file(litellm_params)
return get_cached_authenticator(auth_file) if auth_file else self.authenticator
@staticmethod
def _get_access_token_or_raise(authenticator: Authenticator, model: str, llm_provider: str) -> str:
try:
return authenticator.get_access_token()
except GetAccessTokenError as e:
raise AuthenticationError(
model=model,
llm_provider=llm_provider,
message=str(e),
)
def api_base_without_login(self) -> str:
return self.authenticator.get_api_base()
@ -32,16 +56,16 @@ class ChatGPTConfig(OpenAIConfig):
api_base: str | None,
api_key: str | None,
custom_llm_provider: str,
litellm_params: Mapping[str, object] | GenericLiteLLMParams | None = None,
) -> tuple[str | None, str | None, str]:
dynamic_api_base: Final = self.api_base_without_login()
try:
dynamic_api_key: Final = self.authenticator.get_access_token()
except GetAccessTokenError as e:
raise AuthenticationError(
model=model,
llm_provider=custom_llm_provider,
message=str(e),
)
client_credential: Final = get_chatgpt_request_credential(litellm_params)
if client_credential is not None:
return chatgpt_client_api_base(api_base), client_credential.access_token, custom_llm_provider
if api_key:
return api_base or self.api_base_without_login(), api_key, custom_llm_provider
authenticator: Final = self._resolve_authenticator(litellm_params)
dynamic_api_base: Final = authenticator.get_api_base()
dynamic_api_key: Final = self._get_access_token_or_raise(authenticator, model, custom_llm_provider)
return dynamic_api_base, dynamic_api_key, custom_llm_provider
def get_openai_compatible_provider_info(
@ -50,26 +74,49 @@ class ChatGPTConfig(OpenAIConfig):
api_base: str | None,
api_key: str | None,
custom_llm_provider: str,
litellm_params: Mapping[str, object] | GenericLiteLLMParams | None = None,
) -> tuple[str | None, str | None, str]:
return self._get_openai_compatible_provider_info(model, api_base, api_key, custom_llm_provider)
return self._get_openai_compatible_provider_info(model, api_base, api_key, custom_llm_provider, litellm_params)
def validate_environment(
self,
headers: dict,
headers: dict[str, str],
model: str,
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
litellm_params: dict[str, object],
api_key: str | None = None,
api_base: str | None = None,
) -> dict:
) -> dict[str, str]:
client_credential: Final = get_chatgpt_client_credential(headers)
if client_credential is not None:
return dict(
get_chatgpt_default_headers(
client_credential.access_token,
client_credential.account_id,
ensure_chatgpt_session_id(litellm_params),
),
**strip_chatgpt_client_credential_headers(headers),
)
authenticator: Final = self._resolve_authenticator(litellm_params)
resolved_api_key: Final = (
self._get_access_token_or_raise(authenticator, model, "chatgpt")
if get_chatgpt_auth_file(litellm_params)
else api_key
)
validated_headers: Final = super().validate_environment(
headers, model, messages, optional_params, litellm_params, api_key, api_base
dict(strip_chatgpt_client_credential_headers(headers)),
model,
messages,
optional_params,
litellm_params,
resolved_api_key,
api_base,
)
account_id: Final = self.authenticator.get_account_id()
account_id: Final = authenticator.get_account_id(resolved_api_key)
session_id: Final = ensure_chatgpt_session_id(litellm_params)
default_headers: Final = get_chatgpt_default_headers(api_key or "", account_id, session_id)
default_headers: Final = get_chatgpt_default_headers(resolved_api_key or "", account_id, session_id)
return {**default_headers, **validated_headers}
def post_stream_processing(self, stream: Any) -> ChatGPTToolCallNormalizer:

View file

@ -2,14 +2,32 @@
Constants and helpers for ChatGPT subscription OAuth.
"""
import base64
import json
import os
import platform
from typing import Any, Final
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeAlias
from urllib.parse import urlsplit
from uuid import uuid4
import httpx
from pydantic import JsonValue, TypeAdapter
from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.utils import ProviderSpecificHeader
if TYPE_CHECKING:
from litellm.types.router import GenericLiteLLMParams
_SCOPED_HEADERS: Final[TypeAdapter[ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None]] = TypeAdapter(
ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None
)
_STRING_HEADERS: Final = TypeAdapter(Mapping[str, str])
_PARAMS: Final = TypeAdapter(Mapping[str, object])
# OAuth + API constants (derived from openai/codex)
CHATGPT_AUTH_BASE: Final = "https://auth.openai.com"
@ -19,6 +37,12 @@ CHATGPT_OAUTH_TOKEN_URL: Final = f"{CHATGPT_AUTH_BASE}/oauth/token"
CHATGPT_DEVICE_VERIFY_URL: Final = f"{CHATGPT_AUTH_BASE}/codex/device"
CHATGPT_API_BASE: Final = "https://chatgpt.com/backend-api/codex"
CHATGPT_CLIENT_ID: Final = "app_EMoamEEZ73f0CkXaXp7hrann"
CHATGPT_ACCOUNT_ID_HEADER: Final = "ChatGPT-Account-Id"
OPENAI_AUTH_CLAIM_KEY: Final = "https://api.openai.com/auth"
JsonObject: TypeAlias = Mapping[str, JsonValue]
_JSON_OBJECT_ADAPTER: Final = TypeAdapter(JsonObject)
_EMPTY_CLAIMS: Final[JsonObject] = MappingProxyType({})
DEFAULT_ORIGINATOR: Final = "codex_cli_rs"
DEFAULT_USER_AGENT: Final = "codex_cli_rs/0.0.0 (Unknown 0; unknown) unknown"
@ -229,7 +253,7 @@ def get_chatgpt_default_headers(
access_token: str,
account_id: str | None,
session_id: str | None = None,
) -> dict:
) -> dict[str, str]:
originator: Final = get_chatgpt_originator()
user_agent: Final = get_chatgpt_user_agent(originator)
headers: Final = {
@ -242,30 +266,133 @@ def get_chatgpt_default_headers(
if session_id:
headers["session_id"] = session_id
if account_id:
headers["ChatGPT-Account-Id"] = account_id
headers[CHATGPT_ACCOUNT_ID_HEADER] = account_id
return headers
def decode_jwt_claims(token: str) -> JsonObject:
try:
parts: Final = token.split(".")
if len(parts) < 2:
return _EMPTY_CLAIMS
payload_b64: Final = parts[1] + "=" * (-len(parts[1]) % 4)
payload_bytes: Final = base64.urlsafe_b64decode(payload_b64)
return _JSON_OBJECT_ADAPTER.validate_python(json.loads(payload_bytes.decode("utf-8")))
except Exception:
return _EMPTY_CLAIMS
def extract_chatgpt_account_id(token: str | None) -> str | None:
if not token:
return None
auth_claims: Final = decode_jwt_claims(token).get(OPENAI_AUTH_CLAIM_KEY)
if not isinstance(auth_claims, dict):
return None
account_id: Final = auth_claims.get("chatgpt_account_id")
return account_id if isinstance(account_id, str) and account_id else None
def _bearer_token(value: str) -> str:
scheme, separator, token = value.strip().partition(" ")
if not separator:
return scheme
return token.strip() if scheme.lower() == "bearer" else ""
def is_chatgpt_oauth_key(value: str | None) -> bool:
if value is None:
return False
return OPENAI_AUTH_CLAIM_KEY in decode_jwt_claims(_bearer_token(value))
@dataclass(frozen=True, slots=True)
class ChatGPTClientCredential:
access_token: str
account_id: str | None
def _header_value(headers: Mapping[str, str], name: str) -> str | None:
return next((value for key, value in headers.items() if key.lower() == name), None)
def get_chatgpt_client_credential(headers: Mapping[str, str]) -> ChatGPTClientCredential | None:
auth_header: Final = _header_value(headers, "authorization")
if auth_header is None or not is_chatgpt_oauth_key(auth_header):
return None
access_token: Final = _bearer_token(auth_header)
account_id: Final = _header_value(headers, CHATGPT_ACCOUNT_ID_HEADER.lower()) or extract_chatgpt_account_id(
access_token
)
return ChatGPTClientCredential(access_token=access_token, account_id=account_id)
def select_chatgpt_client_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
if get_chatgpt_client_credential(headers) is None:
return MappingProxyType({})
return MappingProxyType(
{
key: value
for key, value in headers.items()
if (key.lower() == "authorization" and is_chatgpt_oauth_key(value))
or key.lower() == CHATGPT_ACCOUNT_ID_HEADER.lower()
}
)
def strip_chatgpt_client_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
return MappingProxyType(
{
key: value
for key, value in headers.items()
if key.lower() not in ("authorization", CHATGPT_ACCOUNT_ID_HEADER.lower())
}
)
def get_chatgpt_request_credential(
params: "Mapping[str, object] | GenericLiteLLMParams | None",
) -> ChatGPTClientCredential | None:
normalized: Final = _normalize_litellm_params(params)
scoped: Final = normalized.get("provider_specific_header")
extra: Final = normalized.get("extra_headers")
headers: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers(
provider_specific_header=_SCOPED_HEADERS.validate_python(scoped), custom_llm_provider="chatgpt"
)
return get_chatgpt_client_credential(headers) or get_chatgpt_client_credential(
_STRING_HEADERS.validate_python(extra) if extra is not None else MappingProxyType({})
)
def chatgpt_client_api_base(api_base: str | None) -> str:
base: Final = api_base or CHATGPT_API_BASE
parsed: Final = urlsplit(base)
if (
parsed.scheme != "https"
or parsed.hostname != "chatgpt.com"
or parsed.port not in (None, 443)
or parsed.username
or parsed.password
):
raise ChatGPTAuthError(
status_code=400, message="Client ChatGPT credentials require the official HTTPS ChatGPT endpoint"
)
return base
def get_chatgpt_default_instructions() -> str:
return os.getenv("CHATGPT_DEFAULT_INSTRUCTIONS") or CHATGPT_DEFAULT_INSTRUCTIONS
def _normalize_litellm_params(litellm_params: Any | None) -> dict:
if litellm_params is None:
return {}
if isinstance(litellm_params, dict):
return litellm_params
if hasattr(litellm_params, "model_dump"):
try:
return litellm_params.model_dump()
except Exception:
return {}
if hasattr(litellm_params, "dict"):
try:
return litellm_params.dict()
except Exception:
return {}
return {}
def _normalize_litellm_params(litellm_params: object) -> Mapping[str, object]:
if isinstance(litellm_params, Mapping):
return _PARAMS.validate_python(litellm_params)
serializer: Final = getattr(litellm_params, "model_dump", None) or getattr(litellm_params, "dict", None)
if not callable(serializer):
return MappingProxyType({})
try:
return _PARAMS.validate_python(serializer())
except Exception:
return MappingProxyType({})
def get_chatgpt_session_id(litellm_params: object) -> str | None:
@ -274,8 +401,8 @@ def get_chatgpt_session_id(litellm_params: object) -> str | None:
value = params.get(key)
if value:
return str(value)
metadata: Final = params.get("metadata")
if isinstance(metadata, dict):
metadata: Final = _normalize_litellm_params(params.get("metadata"))
if metadata:
value = metadata.get("session_id")
if value:
return str(value)

View file

@ -25,13 +25,15 @@ from litellm.types.llms.openai import (
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
from ..authenticator import Authenticator
from ..authenticator import Authenticator, get_cached_authenticator, get_chatgpt_auth_file
from ..common_utils import (
CHATGPT_API_BASE,
GetAccessTokenError,
ensure_chatgpt_session_id,
get_chatgpt_client_credential,
get_chatgpt_default_headers,
get_chatgpt_default_instructions,
strip_chatgpt_client_credential_headers,
)
if TYPE_CHECKING:
@ -49,14 +51,35 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.CHATGPT
def _resolve_authenticator(self, litellm_params: GenericLiteLLMParams | None) -> Authenticator:
auth_file: Final = get_chatgpt_auth_file(litellm_params)
return get_cached_authenticator(auth_file) if auth_file else self.authenticator
def validate_environment(
self,
headers: dict,
headers: dict[str, str],
model: str,
litellm_params: GenericLiteLLMParams | None,
) -> dict:
) -> dict[str, str]:
session_id: Final = ensure_chatgpt_session_id(litellm_params)
client_credential: Final = get_chatgpt_client_credential(headers)
if client_credential is not None:
return {
**get_chatgpt_default_headers(client_credential.access_token, client_credential.account_id, session_id),
**strip_chatgpt_client_credential_headers(headers),
}
authenticator: Final = self._resolve_authenticator(litellm_params)
safe_headers: Final = strip_chatgpt_client_credential_headers(headers)
configured_api_key: Final = litellm_params.api_key if litellm_params is not None else None
if configured_api_key and not get_chatgpt_auth_file(litellm_params):
return {
**get_chatgpt_default_headers(
configured_api_key, authenticator.get_account_id(configured_api_key), session_id
),
**safe_headers,
}
try:
access_token: Final = self.authenticator.get_access_token()
access_token: Final = authenticator.get_access_token()
except GetAccessTokenError as e:
raise AuthenticationError(
model=model,
@ -64,10 +87,9 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
message=str(e),
)
account_id: Final = self.authenticator.get_account_id()
session_id: Final = ensure_chatgpt_session_id(litellm_params)
account_id: Final = authenticator.get_account_id()
default_headers: Final = get_chatgpt_default_headers(access_token, account_id, session_id)
return {**default_headers, **headers}
return {**default_headers, **safe_headers}
def transform_responses_api_request(
self,
@ -79,7 +101,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
) -> dict:
request: Final = super().transform_responses_api_request(
model,
input,
[{"role": "user", "content": input}] if isinstance(input, str) else input,
response_api_optional_request_params,
litellm_params,
headers,

View file

@ -198,6 +198,26 @@ from litellm.utils import (
)
def _merge_headers_case_insensitive(
headers: Mapping[str, object], extra_headers: Mapping[str, object] | None
) -> dict[str, str]: # mutable-ok: legacy HTTP handler methods require a dict of headers
overridden: Final = frozenset(name.lower() for name in extra_headers or {})
return {
**{name: str(value) for name, value in headers.items() if name.lower() not in overridden},
**{name: str(value) for name, value in (extra_headers or {}).items()},
}
def _client_request_headers(
response_api_optional_request_params: Mapping[str, object],
extra_headers: Mapping[str, object] | None,
) -> dict[str, str]: # mutable-ok: provider validate_environment methods require a dict
from litellm.litellm_core_utils.core_helpers import as_str_mapping
optional_headers: Final = as_str_mapping(response_api_optional_request_params.get("extra_headers"))
return _merge_headers_case_insensitive(optional_headers or {}, extra_headers)
def _rust_responses_websocket_enabled(
custom_llm_provider: str | None,
) -> bool:
@ -2906,14 +2926,18 @@ class BaseLLMHTTPHandler:
else:
sync_httpx_client = client
headers = responses_api_provider_config.validate_environment(
headers=response_api_optional_request_params.get("extra_headers", {}) or {},
model=model,
litellm_params=litellm_params,
)
if extra_headers:
headers.update(extra_headers)
headers = { # rebind-ok: validation replaces the method's incoming headers
**responses_api_provider_config.validate_environment(
headers=(
_client_request_headers(response_api_optional_request_params, extra_headers)
if custom_llm_provider == "chatgpt"
else response_api_optional_request_params.get("extra_headers", {}) or {}
),
model=model,
litellm_params=litellm_params,
),
**((extra_headers or {}) if custom_llm_provider != "chatgpt" else {}),
}
# Check if streaming is requested
stream = response_api_optional_request_params.get("stream", False)
@ -3104,11 +3128,15 @@ class BaseLLMHTTPHandler:
validated_headers: Final = {
**await asyncio.to_thread(
responses_api_provider_config.validate_environment,
headers=response_api_optional_request_params.get("extra_headers", {}) or {},
headers=(
_client_request_headers(response_api_optional_request_params, extra_headers)
if custom_llm_provider == "chatgpt"
else response_api_optional_request_params.get("extra_headers", {}) or {}
),
model=model,
litellm_params=litellm_params,
),
**(extra_headers or {}),
**((extra_headers or {}) if custom_llm_provider != "chatgpt" else {}),
}
# Check if streaming is requested

View file

@ -56,9 +56,9 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig):
# Response format mappings from OpenAI to ElevenLabs
FORMAT_MAPPINGS = {
"mp3": "mp3_44100_128",
"pcm": "pcm_44100",
"pcm": "pcm_24000",
"opus": "opus_48000_128",
# ElevenLabs does not support WAV, AAC, or FLAC formats.
"wav": "wav_24000",
}
ELEVENLABS_QUERY_PARAMS_KEY = "__elevenlabs_query_params__"

View file

@ -726,6 +726,9 @@ def gemini_convert_messages_with_history(
from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig
forward_function_call_id: Final = VertexGeminiConfig.forward_gemini_function_call_id(model or "")
supports_multimodal_function_responses: Final = model is None or (
VertexGeminiConfig.supports_multimodal_function_responses(model)
)
try:
while msg_i < len(messages):
@ -1029,6 +1032,7 @@ def gemini_convert_messages_with_history(
messages[msg_i],
last_message_with_tool_calls,
forward_function_call_id=forward_function_call_id,
supports_multimodal_function_responses=supports_multimodal_function_responses,
)
msg_i += 1
# Handle both single part and list of parts (for Computer Use with images)

View file

@ -689,6 +689,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
) -> bool:
return cls._forward_gemini_function_call_id(model)
@classmethod
def supports_multimodal_function_responses(cls, model: str) -> bool:
return cls.is_gemini_3_or_newer(model)
def _supports_penalty_parameters(self, model: str) -> bool:
# Gemini 3 models do not support penalty parameters
if VertexGeminiConfig._is_gemini_3_or_newer(model):

View file

@ -677,7 +677,11 @@ async def acompletion(
}
if custom_llm_provider is None:
_supplemental: Final = cast( # cast-ok: kwargs values are request-scoped objects
"dict[str, object]", {k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS if k in kwargs}
"dict[str, object]",
{
**{k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS | {"provider_specific_header"} if k in kwargs},
"extra_headers": extra_headers,
},
)
_, custom_llm_provider, _, _ = get_llm_provider(
model=model,
@ -5678,7 +5682,10 @@ def completion(
if deployment_id is not None: # azure llms
model = deployment_id
custom_llm_provider = "azure"
_supplemental_provider_params: Final = {k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS if k in kwargs}
_supplemental_provider_params: Final = {
**{k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS | {"provider_specific_header"} if k in kwargs},
"extra_headers": extra_headers,
}
model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
@ -5998,6 +6005,7 @@ def completion(
tpm=kwargs.get("tpm"),
rpm=kwargs.get("rpm"),
use_xai_oauth=kwargs.get("use_xai_oauth", False),
chatgpt_auth_file=routing_kwargs.get("chatgpt_auth_file"),
gigachat_scope=kwargs.get("gigachat_scope"),
gigachat_auth_url=kwargs.get("gigachat_auth_url"),
gigachat_access_token=kwargs.get("gigachat_access_token"),
@ -9140,6 +9148,7 @@ async def ahealth_check(
custom_llm_provider=custom_llm_provider_from_params,
api_base=api_base_from_params,
api_key=api_key_from_params,
litellm_params=GenericLiteLLMParams(**model_params),
)
model_params["cache"] = {"no-cache": True} # don't used cached responses for making health check calls

View file

@ -6206,6 +6206,8 @@
"text",
"audio"
],
"supports_audio_input": true,
"supports_audio_output": true,
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
@ -6270,6 +6272,8 @@
"text",
"audio"
],
"supports_audio_input": true,
"supports_audio_output": true,
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
@ -33698,7 +33702,7 @@
"input_cost_per_token": 1e-05,
"input_cost_per_token_above_272k_tokens": 2e-05,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -33722,7 +33726,7 @@
"input_cost_per_token": 1e-07,
"input_cost_per_token_above_272k_tokens": 2e-07,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -33746,7 +33750,7 @@
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_272k_tokens": 4e-06,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -33770,7 +33774,7 @@
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_272k_tokens": 4e-06,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -78851,5 +78855,47 @@
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true
},
"baseten/decision-models/baseten-elephant-9b-latest": {
"input_cost_per_token": 0.0,
"litellm_provider": "baseten",
"max_input_tokens": 16384,
"mode": "decisions",
"output_cost_per_token": 0.0,
"source": "https://inference.baseten.co/v1/models",
"supports_vision": true
},
"baseten/decision-models/baseten-mammoth-27b-latest": {
"input_cost_per_token": 0.0,
"litellm_provider": "baseten",
"max_input_tokens": 16384,
"mode": "decisions",
"output_cost_per_token": 0.0,
"source": "https://inference.baseten.co/v1/models",
"supports_vision": true
},
"github_copilot/gpt-6.1-terra": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
"cache_read_input_token_cost": 5e-08,
"cache_read_input_token_cost_above_272k_tokens": 1e-07,
"input_cost_per_token": 1e-06,
"input_cost_per_token_above_272k_tokens": 2e-06,
"litellm_provider": "github_copilot",
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 5e-06,
"output_cost_per_token_above_272k_tokens": 7.5e-06,
"source": "https://github.blog/changelog/2026-10-26-gpt-6-1-terra-in-github-copilot",
"supported_endpoints": [
"/v1/responses"
],
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_vision": true
}
}

View file

@ -27,6 +27,7 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import (
get_passthrough_resource_metadata_url,
get_passthrough_www_authenticate,
get_request_base_url,
get_route_relative_request_path,
well_known_root_suffix,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import (
@ -387,7 +388,7 @@ def _gateway_dcr_challenge(
"""The RFC 9728 challenge pointing the client at the protected-resource metadata
matching the scope it requested: the per-server document (same URL spelling the
request arrived on) when the single target advertises gateway-owned sign-in,
else the gateway's aggregate document. Either way the client discovers the gateway
else a gateway document containing the public resource. Either way the client discovers the gateway
as its authorization server and starts the same sign-in flow.
``invalid_token`` adds the RFC 6750 error code for a request that DID
@ -397,10 +398,13 @@ def _gateway_dcr_challenge(
client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
target: Final = _gateway_dcr_challenge_target(route, mcp_servers, client_ip)
server_scope: Final = gateway_dcr_challenge_scope(route, mcp_servers, client_ip) if target is None else None
public_path: Final = get_route_relative_request_path(request.scope).rstrip("/")
aggregate_resource: Final = public_path == "/mcp" or public_path.startswith("/mcp/")
metadata_suffix: Final = f"{well_known_root_suffix()}/mcp" if aggregate_resource else ""
resource_metadata_url: Final = (
get_passthrough_resource_metadata_url(request.scope, target)
if target is not None
else f"{get_request_base_url(request)}/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp"
else f"{get_request_base_url(request)}/.well-known/oauth-protected-resource{metadata_suffix}"
)
error_attr: Final = 'error="invalid_token", ' if invalid_token else ""
scope_attr: Final = f', scope="{server_scope}"' if server_scope is not None else ""

View file

@ -33,9 +33,8 @@ class OAuthScopeResponse:
self._request_path = request_path
self._ready = asyncio.Event()
self._lock = asyncio.Lock()
self._start: Message | None = None
self._challenge: MCPUpstreamAuthError | None = None
self._refused = False
self._closed = False
@property
def pending(self) -> bool:
@ -50,31 +49,33 @@ class OAuthScopeResponse:
self._ready.set()
async def __call__(self, message: Message) -> None:
if message["type"] == "http.response.start":
self._start = message
if message["status"] != 200:
self.allow()
return
await self._ready.wait()
async with self._lock:
if self._refused:
if self._closed:
return
if self._start is not None:
start: Final = self._start
self._start = None
if self._challenge is not None:
error: Final = self._challenge.to_http_exception(self._base_url, self._request_path)
response: Final = JSONResponse(
status_code=error.status_code, content={"detail": error.detail}, headers=error.headers
)
self._refused = True
await self._send(
{"type": "http.response.start", "status": response.status_code, "headers": response.raw_headers}
)
await self._send({"type": "http.response.body", "body": bytes(response.body)})
return
await self._send(start)
await self._send(message)
try:
if message["type"] == "http.response.start":
if message["status"] != 200:
self.allow()
await self._ready.wait()
if self._challenge is not None:
error: Final = self._challenge.to_http_exception(self._base_url, self._request_path)
response: Final = JSONResponse(
status_code=error.status_code, content={"detail": error.detail}, headers=error.headers
)
self._closed = True
await self._send(
{
"type": "http.response.start",
"status": response.status_code,
"headers": response.raw_headers,
}
)
await self._send({"type": "http.response.body", "body": bytes(response.body)})
return
await self._send(message)
except BaseException:
self._closed = True
raise
def scope_response_from_request(request: object) -> OAuthScopeResponse | None:

View file

@ -207,6 +207,11 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
module_path="litellm.proxy.management_endpoints.config_override_endpoints",
path_prefixes=("/config_overrides",),
),
LazyFeature(
name="codex_passthrough",
module_path="litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints",
path_prefixes=("/backend-api/codex/",),
),
LazyFeature(
name="llm_passthrough",
module_path="litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints",

View file

@ -8673,6 +8673,184 @@
}
}
},
"codex_passthrough": {
"components": {
"schemas": {
"HTTPValidationError": {
"properties": {
"detail": {
"items": {
"$ref": "#/components/schemas/ValidationError"
},
"title": "Detail",
"type": "array"
}
},
"title": "HTTPValidationError",
"type": "object"
},
"ValidationError": {
"properties": {
"ctx": {
"title": "Context",
"type": "object"
},
"input": {
"title": "Input"
},
"loc": {
"items": {
"anyOf": [
{
"type": "string"
},
{
"type": "integer"
}
]
},
"title": "Location",
"type": "array"
},
"msg": {
"title": "Message",
"type": "string"
},
"type": {
"title": "Error Type",
"type": "string"
}
},
"required": [
"loc",
"msg",
"type"
],
"title": "ValidationError",
"type": "object"
}
}
},
"paths": {
"/backend-api/codex/{endpoint}": {
"get": {
"description": "Forward a client's ChatGPT subscription with separate LiteLLM authentication.",
"operationId": "codex_passthrough_route_backend_api_codex__endpoint__get",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"enum": [
"models",
"responses"
],
"title": "Endpoint",
"type": "string"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Codex Passthrough Route",
"tags": [
"codex_passthrough"
]
},
"post": {
"description": "Forward a client's ChatGPT subscription with separate LiteLLM authentication.",
"operationId": "codex_passthrough_route_backend_api_codex__endpoint__post",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"enum": [
"models",
"responses"
],
"title": "Endpoint",
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
"schema": {
"additionalProperties": true,
"type": "object"
}
}
},
"required": true
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"additionalProperties": true,
"type": "object"
}
},
"text/event-stream": {
"schema": {
"type": "string"
}
}
},
"description": "Unmodified Codex response or server-sent events"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Codex Passthrough Route",
"tags": [
"codex_passthrough"
]
}
}
}
},
"compliance": {
"components": {
"schemas": {

View file

@ -391,6 +391,7 @@ class LiteLLMRoutes(enum.Enum):
]
mapped_pass_through_routes = [
"/backend-api/codex",
"/bedrock",
"/comprehendmedical",
"/azure_speech",
@ -1766,8 +1767,10 @@ class BudgetNewRequest(LiteLLMPydanticObjectBase):
max_parallel_requests: int | None = Field(
default=None, description="Max concurrent requests allowed for this budget id."
)
tpm_limit: int | None = Field(default=None, description="Max tokens per minute, allowed for this budget id.")
rpm_limit: int | None = Field(default=None, description="Max requests per minute, allowed for this budget id.")
tpm_limit: int | None = Field(default=None, ge=0, description="Max tokens per minute, allowed for this budget id.")
rpm_limit: int | None = Field(
default=None, ge=0, description="Max requests per minute, allowed for this budget id."
)
tpd_limit: int | None = Field(
default=None, description="Max tokens per day, charged by batch submissions, allowed for this budget id."
)
@ -1878,6 +1881,8 @@ RouterSettingsDict = Annotated[
class NewTeamRequest(TeamBase):
tpm_limit: int | None = Field(default=None, ge=0)
rpm_limit: int | None = Field(default=None, ge=0)
router_settings: RouterSettingsDict | None = None
model_aliases: dict | None = None
model_max_budget: GenericBudgetConfigType | None = Field(
@ -1957,8 +1962,8 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
team_alias: str | None = None
organization_id: str | None = None
metadata: dict | None = None
tpm_limit: int | None = None
rpm_limit: int | None = None
tpm_limit: int | None = Field(default=None, ge=0)
rpm_limit: int | None = Field(default=None, ge=0)
tpd_limit: int | None = None
max_budget: float | None = None
soft_budget: float | None = None

View file

@ -6461,8 +6461,29 @@ def _search_tool_deny_by_default(general_settings: Mapping[str, object]) -> bool
return True
def _is_dashboard_session(valid_token: UserAPIKeyAuth | None) -> bool:
return valid_token is not None and valid_token.team_id == UI_TEAM_ID
def _is_search_tool_strict_identity(valid_token: UserAPIKeyAuth) -> bool:
if not _is_dashboard_session(valid_token):
return _is_strict_grant_identity(valid_token)
return valid_token.user_role not in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
)
def _search_tool_strict_layers(
valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None
) -> frozenset[GrantLayer]:
if _is_dashboard_session(valid_token):
return frozenset(("user",))
return _strict_grant_layers(True, valid_token, team_object)
def is_search_tool_deny_by_default_applied(valid_token: UserAPIKeyAuth, general_settings: Mapping[str, object]) -> bool:
return _search_tool_deny_by_default(general_settings) and _is_strict_grant_identity(valid_token)
return _search_tool_deny_by_default(general_settings) and _is_search_tool_strict_identity(valid_token)
def _search_tool_denied(object_type: GrantLayer, message: str) -> ProxyException:
@ -6506,12 +6527,12 @@ async def resolve_search_tool_grants(
load_team_object: TeamObjectLoader,
) -> SearchToolGrants:
"""
Without `general_settings.search_tool_deny_by_default`, or for the master key and dashboard sessions,
Without `general_settings.search_tool_deny_by_default`, or for the master key and admin dashboard sessions,
the key and team lists only restrict when nonempty. With it, the identities `_strict_grant_layers` names
must each list the tool, and a missing record, `null`, `[]`, or a team or user that fails to load grants nothing
"""
if not is_search_tool_deny_by_default_applied(valid_token, general_settings):
team_object: Final = await load_team_object()
team_object: Final = await load_team_object() if not _is_dashboard_session(valid_token) else None
return SearchToolGrants(
grants=MappingProxyType(
{"key": valid_token.object_permission, "team": team_object.object_permission if team_object else None}
@ -6521,8 +6542,10 @@ async def resolve_search_tool_grants(
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
strict_team_object: Final = await _strict_team_object(load_team_object)
strict_layers: Final = _strict_grant_layers(True, valid_token, strict_team_object)
strict_team_object: Final = (
await _strict_team_object(load_team_object) if not _is_dashboard_session(valid_token) else None
)
strict_layers: Final = _search_tool_strict_layers(valid_token, strict_team_object)
if prisma_client is None:
return SearchToolGrants(grants=MappingProxyType({}), strict_layers=strict_layers)
user_object: Final = (
@ -6607,7 +6630,7 @@ def check_unregistered_search_fallback(
"""The unregistered provider fallback names no search tool that a grant could list, so deny-by-default denies it"""
if not is_search_tool_deny_by_default_applied(valid_token, general_settings):
return True
layer: Final = min(_strict_grant_layers(True, valid_token, None), key=("key", "team", "user").index)
layer: Final = min(_search_tool_strict_layers(valid_token, None), key=("key", "team", "user").index)
raise _search_tool_denied(
layer,
"No registered search tool is available and search_tool_deny_by_default is enabled",

View file

@ -736,7 +736,9 @@ def is_request_body_safe(
def _reject_request_credential_params(litellm_params: Mapping[str, object], *, manages_deployments: bool) -> None:
if manages_deployments:
reject_server_owned_wif_fields(litellm_params)
reject_server_owned_wif_fields(
MappingProxyType({key: value for key, value in litellm_params.items() if key != "chatgpt_auth_file"})
)
else:
reject_server_owned_wif_params(litellm_params)
reject_federated_credential_reference(litellm_params)

View file

@ -119,6 +119,7 @@ from litellm.proxy.common_utils.sse_keepalive import (
)
from litellm.proxy.dd_span_tagger import DDSpanTagger
from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression
from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash
from litellm.proxy.native_compaction import with_proxy_compaction_executor
from litellm.proxy.route_llm_request import (
route_request,
@ -1416,6 +1417,8 @@ async def open_sse_before_first_byte(
if interval is None:
return await produce_response
# Share admission state with stream cleanup before the producer copies this context.
get_or_create_request_stash()
produce_task: Final = asyncio.ensure_future(produce_response)
await asyncio.wait((produce_task,), timeout=interval)
if produce_task.done():

View file

@ -2304,7 +2304,15 @@ async def test_model_connection(
)
selected_by_id: Final = deployment_by_id is not None
reject_server_owned_wif_fields(request_litellm_params)
reject_server_owned_wif_fields(
MappingProxyType(
{
key: value
for key, value in _OBJECT_MAPPING.validate_python(request_litellm_params).items()
if key != "chatgpt_auth_file"
}
)
)
# Merge: config params (from proxy config) as base, request params override
litellm_params = {
**_config_base_for_health_check(

View file

@ -1417,13 +1417,14 @@ def clean_headers(
Cleaned headers dict
"""
from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key
from litellm.llms.chatgpt.common_utils import is_chatgpt_oauth_key
clean_headers: Final = {}
litellm_key_lower: Final = litellm_key_header_name.lower() if litellm_key_header_name is not None else None
for header, value in headers.items():
header_lower = header.lower()
if header_lower == "authorization" and is_anthropic_oauth_key(value):
if header_lower == "authorization" and (is_anthropic_oauth_key(value) or is_chatgpt_oauth_key(value)):
if authenticated_with_header is None or authenticated_with_header.lower() != "authorization":
clean_headers[header] = value
continue
@ -1452,7 +1453,7 @@ def clean_headers(
def _is_credential_header(header: str) -> bool:
"""Whether `header` carries a caller credential rather than request context."""
return header.lower() in _CREDENTIAL_HEADER_NAMES
return header.lower().removeprefix("x-pass-") in _CREDENTIAL_HEADER_NAMES
def redact_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
@ -4282,6 +4283,7 @@ _NATIVE_MESSAGES_ANTHROPIC_API_HEADER_PROVIDERS: Final = ",".join(
(LlmProviders.AZURE_AI.value, LlmProviders.DATABRICKS.value)
)
_ANTHROPIC_OAUTH_CREDENTIAL_PROVIDERS: Final = ",".join(sorted(ANTHROPIC_OAUTH_FORWARD_PROVIDERS))
_CHATGPT_OAUTH_CREDENTIAL_PROVIDERS: Final = LlmProviders.CHATGPT.value
def add_provider_specific_headers_to_request(
@ -4289,6 +4291,7 @@ def add_provider_specific_headers_to_request(
headers: Mapping[str, str],
) -> bool:
from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key
from litellm.llms.chatgpt.common_utils import select_chatgpt_client_credential_headers
anthropic_api_headers: Final = {header: headers[header] for header in ANTHROPIC_API_HEADERS if header in headers}
anthropic_oauth_credential_headers: Final = {
@ -4296,6 +4299,8 @@ def add_provider_specific_headers_to_request(
for header, value in headers.items()
if header.lower() == "authorization" and is_anthropic_oauth_key(value)
}
chatgpt_oauth_credential_headers: Final = dict(select_chatgpt_client_credential_headers(headers))
typesafe_danger_zone_headers: Final = {
header: value
for header, value in headers.items()
@ -4320,13 +4325,17 @@ def add_provider_specific_headers_to_request(
ProviderSpecificHeader(
custom_llm_provider=LlmProviders.TYPESAFE.value, extra_headers=typesafe_danger_zone_headers
),
ProviderSpecificHeader(
custom_llm_provider=_CHATGPT_OAUTH_CREDENTIAL_PROVIDERS,
extra_headers=chatgpt_oauth_credential_headers,
),
)
if scoped_header["extra_headers"]
]
if scoped_headers:
data["provider_specific_header"] = scoped_headers[0] if len(scoped_headers) == 1 else scoped_headers
return bool(anthropic_oauth_credential_headers)
return bool(anthropic_oauth_credential_headers or chatgpt_oauth_credential_headers)
def _add_otel_traceparent_to_data(data: dict, request: Request):

View file

@ -0,0 +1,87 @@
import hmac
from types import MappingProxyType
from typing import Annotated, Final, Literal
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from litellm.llms.chatgpt.common_utils import CHATGPT_API_BASE
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import pass_through_request
router: Final = APIRouter()
_HEADERS: Final = frozenset(
{
"authorization",
"chatgpt-account-id",
"content-type",
"accept",
"user-agent",
"originator",
"openai-beta",
"session-id",
"thread-id",
"session_id",
"conversation_id",
"x-codex-turn-metadata",
"x-codex-turn-state",
}
)
@router.get("/backend-api/codex/{endpoint}", tags=["ChatGPT Pass-through", "pass-through"])
@router.post(
"/backend-api/codex/{endpoint}",
tags=["ChatGPT Pass-through", "pass-through"],
openapi_extra={
"requestBody": {
"required": True,
"content": {"application/json": {"schema": {"type": "object", "additionalProperties": True}}},
}
},
responses={
200: {
"description": "Unmodified Codex response or server-sent events",
"content": {
"application/json": {"schema": {"type": "object", "additionalProperties": True}},
"text/event-stream": {"schema": {"type": "string"}},
},
}
},
)
async def codex_passthrough_route(
endpoint: Literal["models", "responses"],
request: Request,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> Response:
"""Forward a client's ChatGPT subscription with separate LiteLLM authentication."""
expected_method: Final = "GET" if endpoint == "models" else "POST"
if request.method != expected_method:
raise HTTPException(status_code=405, detail="Method not allowed", headers={"Allow": expected_method})
authorization: Final = request.headers.get("authorization", "")
scheme, _, token = authorization.partition(" ")
gateway_key: Final = request.headers.get("x-litellm-api-key", "")
gateway_parts: Final = gateway_key.split()
gateway_token: Final = (
gateway_parts[-1]
if len(gateway_parts) == 1 or (len(gateway_parts) == 2 and gateway_parts[0].lower() == "bearer")
else ""
)
if scheme.lower() != "bearer" or not token.strip() or not gateway_token:
raise HTTPException(
status_code=401,
detail="Send the ChatGPT bearer in Authorization and the LiteLLM key in x-litellm-api-key",
)
if hmac.compare_digest(token.strip().encode(), gateway_token.encode()):
raise HTTPException(status_code=401, detail="ChatGPT and LiteLLM credentials must be separate")
headers: Final = MappingProxyType({name: value for name, value in request.headers.items() if name in _HEADERS})
return await pass_through_request(
request=request,
target=f"{CHATGPT_API_BASE}/{endpoint}",
custom_headers=dict(headers),
user_api_key_dict=user_api_key_dict,
custom_llm_provider="chatgpt",
upstream_headers_filter=lambda _: headers,
)

View file

@ -10,6 +10,7 @@ from typing import Final, Literal, TypeAlias
from urllib.parse import urlparse
import httpx
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -95,6 +96,8 @@ def is_openai_compatible_url(url_route: str | None) -> bool:
hostname: Final = parsed_url.hostname
if not hostname:
return False
if hostname == "chatgpt.com" and parsed_url.path.startswith("/backend-api/codex/"):
return True
if _hostname_matches(hostname, _OPENAI_HOSTNAMES):
return True
if _hostname_matches(hostname, _AZURE_OPENAI_HOSTNAMES):
@ -198,7 +201,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
if not url_route:
return False
parsed_url: Final = urlparse(url_route)
return _is_openai_compatible_host(parsed_url.hostname) and (
return is_openai_compatible_url(url_route) and (
"/v1/responses" in parsed_url.path or "/responses" in parsed_url.path
)
@ -424,7 +427,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
) = None
handler_instance: Final = OpenAIPassthroughLoggingHandler()
custom_llm_provider: Final = kwargs.get("custom_llm_provider", "openai")
custom_llm_provider: Final = TypeAdapter(str).validate_python(
kwargs.get("custom_llm_provider", logging_obj.model_call_details.get("custom_llm_provider", "openai"))
)
if is_chat_completions:
# Handle chat completions with existing logic

View file

@ -118,6 +118,7 @@ from litellm.proxy.litellm_pre_call_utils import ( # noqa: F401 # legacy modul
_key_or_team_allows_client_pricing_override, # pyright: ignore[reportPrivateUsage] # reuse the proxy's pricing trust policy
_strip_client_pricing_overrides, # pyright: ignore[reportPrivateUsage] # sanitize before trusted hooks add guardrail costs
get_dynamic_logging_metadata, # pyright: ignore[reportPrivateUsage] # shared proxy helper, same import style as _read_request_body above
redact_credential_headers,
)
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
from litellm.proxy.route_priority import (
@ -409,6 +410,8 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
@staticmethod
def get_endpoint_type(url: str, custom_llm_provider: str | None = None) -> EndpointType:
parsed_url: Final = urlparse(url)
if custom_llm_provider == "chatgpt":
return EndpointType.OPENAI
if custom_llm_provider in ("typesafe", "laya", "bespoke") and parsed_url.path.removesuffix("/").endswith(
"/v1/systemone"
):
@ -723,7 +726,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
"url": str(request.url),
"method": request.method,
"body": copy.copy(_parsed_body), # use copy instead of deepcopy
"headers": request.headers,
"headers": redact_credential_headers(request.headers),
},
},
"call_type": "pass_through_endpoint",
@ -1490,7 +1493,7 @@ async def pass_through_request(
additional_args={
"complete_input_dict": _parsed_body,
"api_base": str(logging_url),
"headers": upstream_headers,
"headers": redact_credential_headers(upstream_headers),
},
)
stream = HttpPassThroughEndpointHelpers.update_stream_param_based_on_request_body(
@ -1561,7 +1564,7 @@ async def pass_through_request(
data=_parsed_body or {},
user_api_key_dict=user_api_key_dict,
response=relay_response,
request_headers=dict(request.headers),
request_headers=dict(redact_credential_headers(request.headers)),
)
if callback_headers:
_response_headers.update(callback_headers)
@ -1653,7 +1656,7 @@ async def pass_through_request(
data=_parsed_body or {},
user_api_key_dict=user_api_key_dict,
response=detected_relay_response,
request_headers=dict(request.headers),
request_headers=dict(redact_credential_headers(request.headers)),
)
if callback_headers:
_response_headers.update(callback_headers)
@ -1703,7 +1706,7 @@ async def pass_through_request(
data=_parsed_body or {},
user_api_key_dict=user_api_key_dict,
response=response,
request_headers=dict(request.headers),
request_headers=dict(redact_credential_headers(request.headers)),
)
if relay_callback_headers:
relay_custom_headers.update(relay_callback_headers)
@ -1878,7 +1881,7 @@ async def pass_through_request(
data=_parsed_body or {},
user_api_key_dict=user_api_key_dict,
response=response,
request_headers=dict(request.headers),
request_headers=dict(redact_credential_headers(request.headers)),
)
if callback_headers:
custom_headers.update(callback_headers)

View file

@ -21795,8 +21795,7 @@ async def toolset_mcp_route(toolset_name: str, request: Request):
detail=f"Toolset '{toolset_name}' not found",
)
scope: Final = dict(request.scope)
scope["path"] = "/mcp"
scope: Final = {**request.scope, "_original_path": request.scope.get("path", ""), "path": "/mcp"}
token: Final = mcp_active_toolset_id.set(toolset.toolset_id)
try:

View file

@ -308,10 +308,7 @@ def should_use_polling_for_request(
# Otherwise, check ALL deployments for this model_name in router
elif llm_router is not None:
try:
# Get all deployment indices for this model name
indices: Final = llm_router.model_name_to_deployment_indices.get(model, [])
for idx in indices:
deployment_dict = llm_router.model_list[idx]
for deployment_dict in llm_router.deployments_named(model) or ():
litellm_params = deployment_dict.get("litellm_params", {})
# Check custom_llm_provider first

View file

@ -19,6 +19,7 @@ from pydantic import ConfigDict, Field
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata
from litellm.litellm_core_utils.llm_cost_calc.utils import (
calculate_prompt_caching_savings,
generic_cost_per_token,
@ -69,7 +70,9 @@ def _resolve_model(model: str | None, custom_llm_provider: str | None) -> _Model
if not model:
return None
try:
resolved_model, provider, _, _ = litellm.get_llm_provider(model=model, custom_llm_provider=custom_llm_provider)
resolved_model, provider = resolve_model_and_provider_for_metadata(
model=model, custom_llm_provider=custom_llm_provider
)
except Exception as e: # noqa: BLE001 # get_llm_provider raises for unroutable names; degrade to an unavailable estimate
verbose_proxy_logger.debug(
"savings: cannot resolve provider for model=%s custom_llm_provider=%s (%s)", model, custom_llm_provider, e

View file

@ -705,10 +705,12 @@ def _api_base_kwarg(kwargs: Mapping[str, object]) -> str | None:
return api_base if isinstance(api_base, str) else None
def _dispatched_model_name(model: str, custom_llm_provider: str, api_base: str | None) -> str:
provider_model, _, _, _ = litellm.get_llm_provider(
model=model, custom_llm_provider=custom_llm_provider, api_base=api_base
)
def _dispatched_model_name(
model: str, custom_llm_provider: str, api_base: str | None, litellm_params: GenericLiteLLMParams | None = None
) -> str:
from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata
provider_model, _ = resolve_model_and_provider_for_metadata(model=model, custom_llm_provider=custom_llm_provider)
return _strip_responses_routing_prefix(provider_model)
@ -718,6 +720,7 @@ def _will_bridge_to_chat_completions(
use_chat_completions_api: bool,
model_info: object,
api_base: str | None,
litellm_params: GenericLiteLLMParams | None = None,
) -> bool:
"""``_bridges_to_chat_completions`` for callers running before the provider config is resolved.
@ -734,7 +737,7 @@ def _will_bridge_to_chat_completions(
return True
return _bridges_to_chat_completions(
_resolve_responses_api_provider_config(
_dispatched_model_name(normalized_model[0], custom_llm_provider, api_base),
_dispatched_model_name(normalized_model[0], custom_llm_provider, api_base, litellm_params),
custom_llm_provider,
model_info,
api_base,
@ -750,12 +753,15 @@ def provider_stores_the_conversation(model: str, custom_llm_provider: str | None
first MCP round runs as an internal call that writes no spend log at all.
"""
api_base: Final = _api_base_kwarg(kwargs)
litellm_params: Final = GenericLiteLLMParams(**kwargs)
return not _will_bridge_to_chat_completions(
model,
custom_llm_provider or litellm.get_llm_provider(model=model, api_base=api_base)[1],
custom_llm_provider
or litellm.get_llm_provider(model=model, api_base=api_base, litellm_params=litellm_params)[1],
bool(kwargs.get("use_chat_completions_api")),
kwargs.get("model_info"),
api_base,
litellm_params,
)
@ -905,7 +911,8 @@ async def aresponses(
model=model,
api_base=local_vars.get("base_url", None),
litellm_params=GenericLiteLLMParams(
**cast("dict[str, object]", kwargs) # cast-ok: kwargs is the untyped request dict
extra_headers=extra_headers,
**cast("dict[str, object]", kwargs), # cast-ok: kwargs is the untyped request dict
),
)
# Update local_vars with detected provider (fixes #19782)
@ -936,6 +943,7 @@ async def aresponses(
bool(kwargs.get("use_chat_completions_api")),
kwargs.get("model_info"),
_api_base_kwarg(kwargs),
GenericLiteLLMParams(**kwargs),
),
):
_stamp_injection_points_for_the_hook(kwargs, model, custom_llm_provider, input, tools)
@ -967,7 +975,7 @@ async def aresponses(
original_model=original_model,
swapped_model=model,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
kwargs={**kwargs, "extra_headers": extra_headers},
prompt_id=prompt_id,
)
kwargs.pop("prompt_id", None)
@ -1051,11 +1059,14 @@ def _resolve_prompt_swapped_provider(
kwargs: Mapping[str, object],
prompt_id: str | None,
) -> str:
swapped_provider: Final = litellm.get_llm_provider(model=swapped_model)[1]
litellm_params: Final = GenericLiteLLMParams(**kwargs)
swapped_provider: Final = litellm.get_llm_provider(model=swapped_model, litellm_params=litellm_params)[1]
if kwargs.get("api_key") is None and kwargs.get("api_base") is None:
return swapped_provider
try:
original_provider: Final = custom_llm_provider or litellm.get_llm_provider(model=original_model)[1]
original_provider: Final = (
custom_llm_provider or litellm.get_llm_provider(model=original_model, litellm_params=litellm_params)[1]
)
except litellm.BadRequestError:
return swapped_provider
if swapped_provider == original_provider:
@ -1108,6 +1119,7 @@ def _apply_prompt_management_to_responses_call(
use_chat_completions_api,
kwargs.get("model_info"),
_api_base_kwarg(kwargs),
GenericLiteLLMParams(**kwargs),
),
):
_stamp_injection_points_for_the_hook(kwargs, model, custom_llm_provider, input, tools)
@ -1142,7 +1154,7 @@ def _apply_prompt_management_to_responses_call(
original_model=original_model,
swapped_model=model,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
kwargs={**kwargs, "extra_headers": local_vars.get("extra_headers")},
prompt_id=prompt_id,
)
)
@ -1498,7 +1510,8 @@ def responses(
model=model,
api_base=local_vars.get("base_url", None),
litellm_params=GenericLiteLLMParams(
**cast("dict[str, object]", kwargs) # cast-ok: kwargs is the untyped request dict
extra_headers=extra_headers,
**cast("dict[str, object]", kwargs), # cast-ok: kwargs is the untyped request dict
),
)
local_vars["custom_llm_provider"] = custom_llm_provider
@ -1514,8 +1527,16 @@ def responses(
use_chat_completions_api=use_chat_completions_api,
)
# get llm provider logic
litellm_params: Final = GenericLiteLLMParams(**kwargs)
from litellm.litellm_core_utils.core_helpers import attribute_or_none
litellm_params: Final = GenericLiteLLMParams(extra_headers=extra_headers, **kwargs)
provider_headers: Final = ResponsesAPIRequestUtils.merge_provider_specific_headers(
extra_headers=extra_headers,
provider_specific_header=attribute_or_none(litellm_params, "provider_specific_header"),
custom_llm_provider=custom_llm_provider,
)
local_vars["extra_headers"] = provider_headers
affinity_api_base: Final = (
get_openai_credentials(api_base=litellm_params.api_base).api_base
if custom_llm_provider == "openai"
@ -1525,7 +1546,7 @@ def responses(
try:
effective_extra_headers: Final = (
add_provider_affinity_header(
headers=extra_headers or MappingProxyType({}),
headers=provider_headers or MappingProxyType({}),
litellm_params=MappingProxyType(
{
"provider_affinity_header": litellm_params.provider_affinity_header,
@ -1538,7 +1559,7 @@ def responses(
default_header=default_affinity_header,
)
if litellm_params.provider_affinity_header is not None or default_affinity_header is not None
else extra_headers
else provider_headers
)
except ValueError as affinity_error:
raise litellm.BadRequestError(

View file

@ -12,6 +12,7 @@ import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, iter_items
from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value, is_nested_path
from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig, ReasoningSummaryLevel
from litellm.types.llms.openai import (
AllMessageValues,
@ -29,11 +30,16 @@ from litellm.types.utils import (
CacheCreationTokenDetails,
CompletionTokensDetailsWrapper,
PromptTokensDetailsWrapper,
ProviderSpecificHeader,
SpecialEnums,
Usage,
text_tokens_without_nested_reasoning,
)
_PROVIDER_SPECIFIC_HEADER_ADAPTER: Final[
TypeAdapter[ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None]
] = TypeAdapter(ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None)
class _UsageExtraFields(ExtraItemsTypedDict, total=False, extra_items=object):
pass
@ -233,6 +239,20 @@ class ResponsesAPIRequestUtils:
**extra_headers,
}
@staticmethod
def merge_provider_specific_headers(
extra_headers: Mapping[str, object] | None,
provider_specific_header: object,
custom_llm_provider: str | None,
) -> dict[str, object] | None: # mutable-ok: legacy Responses request fields require a dict
scoped_headers: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers(
provider_specific_header=_PROVIDER_SPECIFIC_HEADER_ADAPTER.validate_python(provider_specific_header),
custom_llm_provider=custom_llm_provider,
)
if extra_headers is None and not scoped_headers:
return None
return {**(extra_headers or {}), **scoped_headers}
@staticmethod
def _check_valid_arg(
supported_params: list[str] | None,

View file

@ -98,7 +98,9 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor, r
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.get_llm_provider_logic import (
declared_authenticating_provider,
inferred_provider,
is_registered_custom_provider,
resolve_model_and_provider_for_metadata,
)
from litellm.litellm_core_utils.hidden_params import (
HIDDEN_PARAMS_ATTR as _HIDDEN_PARAMS_ATTR,
@ -561,6 +563,7 @@ _ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset(
{"model", "api_base", "api_key", "api_version", "responses_api_version"}
)
_ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params"
_FORWARDED_CREDENTIAL_ISOLATION_PROVIDERS: Final = frozenset({"chatgpt"})
_CLAUDE_CODE_SESSION_ID_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$")
_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600
CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_binding"
@ -1440,6 +1443,9 @@ class Router:
# Maps (team_id, team_public_model_name) -> list of indices in model_list
self.team_model_to_deployment_indices: dict[tuple[str, str], list[int]] = {}
self.team_public_model_names: frozenset[str] = frozenset()
# Guards model_list and the three index maps above: writers change them, and readers turn an
# index into a deployment, under it, so a lookup never lands between a removal and its re-index
self._deployment_index_lock = threading.RLock()
# Initialize cache attributes that ``_invalidate_model_group_info_cache``
# and ``_invalidate_access_groups_cache`` touch *before* the first
@ -2180,8 +2186,8 @@ class Router:
public_model_name: Final = deployment.model_info.team_public_model_name
if team_id is None or public_model_name is None:
return False
sibling_indices: Final = self.team_model_to_deployment_indices.get((team_id, public_model_name)) or ()
routable_siblings: Final = self._filter_blocked_deployments([self.model_list[idx] for idx in sibling_indices])
siblings: Final = self._team_deployments_named(team_id, public_model_name) or ()
routable_siblings: Final = self._filter_blocked_deployments([dict(sibling) for sibling in siblings])
return len(routable_siblings) > 1
_OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY})
@ -4798,9 +4804,39 @@ class Router:
deployment_litellm_model_name = deployment["litellm_params"]["model"]
deployment_api_base = deployment["litellm_params"].get("api_base")
deployment_model_name: Final = deployment["model_name"]
if is_clientside_credential(request_kwargs=kwargs):
from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils
request_context: Final = cast( # cast-ok: the router's legacy request dict has string keys
Mapping[str, object], kwargs
)
raw_deployment_params: Final[object] = deployment["litellm_params"]
deployment_context: Final = _LITELLM_PARAMS_ADAPTER.validate_python(
raw_deployment_params.model_dump(exclude_none=True)
if isinstance(raw_deployment_params, LiteLLM_Params)
else raw_deployment_params
)
declared_provider: Final = deployment_context.get("custom_llm_provider")
declared_model: Final = deployment_context.get("model")
deployment_provider: Final = (
declared_provider
if isinstance(declared_provider, str)
else inferred_provider(declared_model if isinstance(declared_model, str) else None)
)
forwarded_authorization: Final = (
ProviderSpecificHeaderUtils.forwarded_authorization(
request_context.get("provider_specific_header"), deployment_provider
)
if deployment_provider in _FORWARDED_CREDENTIAL_ISOLATION_PROVIDERS
else None
)
if is_clientside_credential(request_kwargs=kwargs) or forwarded_authorization is not None:
deployment_pydantic_obj: Final = self._handle_clientside_credential(
deployment=deployment, kwargs=kwargs, function_name=function_name
deployment=deployment,
kwargs={
**kwargs,
**({"api_key": forwarded_authorization} if forwarded_authorization is not None else {}),
},
function_name=function_name,
)
model_info = deployment_pydantic_obj.model_info.model_dump()
deployment_litellm_model_name = deployment_pydantic_obj.litellm_params.model
@ -10827,20 +10863,17 @@ class Router:
# O(k) via the name→indices map: only touch deployments whose name
# is listed in `available_models`, instead of scanning model_list.
for name in config.available_models:
indices = self.model_name_to_deployment_indices.get(name, [])
if not indices:
deployments = self.deployments_named(name)
if not deployments:
continue
d = (self.model_list or [])[indices[0]]
mi = d.get("model_info") if isinstance(d, dict) else d.model_info
mi_dict: dict[str, Any] = mi if isinstance(mi, dict) else (mi.model_dump() if mi else {})
d = deployments[0]
mi_dict: dict[str, Any] = d.get("model_info") or {}
prefs_raw = mi_dict.get("adaptive_router_preferences")
if prefs_raw is not None:
model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw)
# model_info is the conventional pricing location elsewhere in LiteLLM; litellm_params wins if set.
lp = d.get("litellm_params") if isinstance(d, dict) else d.litellm_params
lp_dict: dict[str, Any] = lp if isinstance(lp, dict) else (lp.model_dump() if lp else {})
cost = lp_dict.get("input_cost_per_token")
cost = d.get("litellm_params", {}).get("input_cost_per_token")
if cost is None:
cost = mi_dict.get("input_cost_per_token")
if cost is not None:
@ -10932,10 +10965,11 @@ class Router:
def set_model_list(self, model_list: list):
original_model_list: Final = copy.deepcopy(model_list)
self._discovered_model_info_cache.flush_cache()
self.model_list = []
self.model_id_to_deployment_index_map = {} # Reset the index
self.model_name_to_deployment_indices = {} # Reset the model_name index
self.team_model_to_deployment_indices = {} # Reset the team_model index
with self._deployment_index_lock:
self.model_list = []
self.model_id_to_deployment_index_map = {} # Reset the index
self.model_name_to_deployment_indices = {} # Reset the model_name index
self.team_model_to_deployment_indices = {} # Reset the team_model index
self.team_pattern_routers = {}
self.team_public_model_names = frozenset()
# Reset per-strategy router registries so hot-reload doesn't leave
@ -11029,6 +11063,9 @@ class Router:
litellm_params.model.split("/")[0] in litellm._known_custom_logger_compatible_callbacks
):
return litellm_params.model, None, None, None
if declared_authenticating_provider(litellm_params.model, litellm_params.custom_llm_provider) is not None:
model, provider = Router._model_group_llm_provider(litellm_params)
return model, provider, None, None
resolved: Final = litellm.get_llm_provider(
model=litellm_params.model,
custom_llm_provider=litellm_params.get("custom_llm_provider", None),
@ -11315,26 +11352,27 @@ class Router:
- model: dict - the model to add to the list
- model_id: Optional[str] - the model ID to use for indexing. If None, will try to get from model["model_info"]["id"]
"""
idx: Final = len(self.model_list)
self.model_list.append(model)
self._invalidate_model_group_info_cache()
self._invalidate_access_groups_cache()
with self._deployment_index_lock:
idx: Final = len(self.model_list)
self.model_list.append(model)
self._invalidate_model_group_info_cache()
self._invalidate_access_groups_cache()
# Update model_id index for O(1) lookup
if model_id is not None:
self.model_id_to_deployment_index_map[model_id] = idx
elif model.get("model_info", {}).get("id") is not None:
self.model_id_to_deployment_index_map[model["model_info"]["id"]] = idx
# Update model_id index for O(1) lookup
if model_id is not None:
self.model_id_to_deployment_index_map[model_id] = idx
elif model.get("model_info", {}).get("id") is not None:
self.model_id_to_deployment_index_map[model["model_info"]["id"]] = idx
# Update model_name index for O(1) lookup
model_name: Final = model.get("model_name")
if model_name:
if model_name not in self.model_name_to_deployment_indices:
self.model_name_to_deployment_indices[model_name] = []
self.model_name_to_deployment_indices[model_name].append(idx)
# Update model_name index for O(1) lookup
model_name: Final = model.get("model_name")
if model_name:
if model_name not in self.model_name_to_deployment_indices:
self.model_name_to_deployment_indices[model_name] = []
self.model_name_to_deployment_indices[model_name].append(idx)
# Update team_model index for O(1) team-scoped lookup
self._update_team_model_index(model, idx)
# Update team_model index for O(1) team-scoped lookup
self._update_team_model_index(model, idx)
def upsert_deployment(self, deployment: Deployment) -> Deployment | None:
"""
@ -11361,18 +11399,15 @@ class Router:
# if there is a new litellm param -> then update the deployment
# remove the previous deployment
removal_idx: int | None = None
deployment_id: Final = deployment.model_info.id
deployment_fast_mapping: Final = self.model_id_to_deployment_index_map
if deployment_id in deployment_fast_mapping:
removal_idx = deployment_fast_mapping[deployment_id]
with self._deployment_index_lock:
removal_idx: Final = self.model_id_to_deployment_index_map.get(_deployment_model_id)
if removal_idx is not None:
self.model_list.pop(removal_idx)
self._invalidate_model_group_info_cache()
self._invalidate_access_groups_cache()
self._update_deployment_indices_after_removal(model_id=deployment_id, removal_idx=removal_idx)
self._update_deployment_indices_after_removal(
model_id=_deployment_model_id, removal_idx=removal_idx
)
# Free the outgoing deployment's pre-routing strategy slot (keyed by the
# OLD model_name/tags) before the re-add below re-registers it.
@ -11606,35 +11641,32 @@ class Router:
- The deleted deployment
- OR None (if deleted deployment not found)
"""
deployment_idx = None
if id in self.model_id_to_deployment_index_map:
deployment_idx = self.model_id_to_deployment_index_map[id]
try:
if deployment_idx is not None:
with self._deployment_index_lock:
deployment_idx: Final = self.model_id_to_deployment_index_map.get(id)
if deployment_idx is None:
return None
# Pop the item from the list first
item: Final = self.model_list.pop(deployment_idx)
self._invalidate_model_group_info_cache()
self._invalidate_access_groups_cache()
self._update_deployment_indices_after_removal(model_id=id, removal_idx=deployment_idx)
_budget_limiter: Final = self.get_router_deployment_budget_limiter()
if _budget_limiter is not None:
_budget_limiter.unregister_deployment_budget(model_id=id)
try:
self._unregister_pre_routing_strategy_for_deployment(
deployment=item if isinstance(item, Deployment) else Deployment(**item)
)
except Exception:
verbose_router_logger.exception(
"delete_deployment: could not release pre-routing strategies for model_id=%s; "
"the deployment is out of the model_list and its indices are repaired",
id,
)
return item
else:
return None
_budget_limiter: Final = self.get_router_deployment_budget_limiter()
if _budget_limiter is not None:
_budget_limiter.unregister_deployment_budget(model_id=id)
try:
self._unregister_pre_routing_strategy_for_deployment(
deployment=item if isinstance(item, Deployment) else Deployment(**item)
)
except Exception:
verbose_router_logger.exception(
"delete_deployment: could not release pre-routing strategies for model_id=%s; "
"the deployment is out of the model_list and its indices are repaired",
id,
)
except Exception:
return None
return item
def get_router_deployment_budget_limiter(
self,
@ -11683,6 +11715,18 @@ class Router:
if _budget_limiter is not None:
_budget_limiter.register_deployment_budget(deployment=deployment.to_json(exclude_none=True))
def deployments_named(self, model_name: str) -> tuple[DeploymentTypedDict, ...] | None:
"""The deployments indexed under `model_name`, read under the index lock, or None when the name is not indexed."""
with self._deployment_index_lock:
indices: Final = self.model_name_to_deployment_indices.get(model_name)
return None if indices is None else tuple(self.model_list[idx] for idx in indices)
def _team_deployments_named(self, team_id: str, public_model_name: str) -> tuple[DeploymentTypedDict, ...] | None:
"""The deployments `team_id` serves under `public_model_name`, or None when that pair is not indexed."""
with self._deployment_index_lock:
indices: Final = self.team_model_to_deployment_indices.get((team_id, public_model_name))
return None if indices is None else tuple(self.model_list[idx] for idx in indices)
def get_deployment(self, model_id: str) -> Deployment | None:
"""
Returns -> Deployment or None
@ -11690,9 +11734,10 @@ class Router:
Raise Exception -> if model found in invalid format
"""
# Use O(1) lookup via model_id_to_deployment_index_map only
if model_id in self.model_id_to_deployment_index_map:
idx: Final = self.model_id_to_deployment_index_map[model_id]
model: Final = self.model_list[idx]
with self._deployment_index_lock:
idx: Final = self.model_id_to_deployment_index_map.get(model_id)
model: Final = None if idx is None else self.model_list[idx]
if model is not None:
if isinstance(model, dict):
return Deployment(**model)
elif isinstance(model, Deployment):
@ -11726,17 +11771,17 @@ class Router:
Optimized with O(1) index lookup instead of O(n) linear scan.
"""
# O(1) lookup in model_name index
if model_group_name in self.model_name_to_deployment_indices:
indices: Final = self.model_name_to_deployment_indices[model_group_name]
if indices:
# Return first deployment for this model_name
model: Final = self.model_list[indices[0]]
if isinstance(model, dict):
return Deployment(**model)
elif isinstance(model, Deployment):
return model
else:
raise Exception(f"Model Name invalid - {type(model)}")
with self._deployment_index_lock:
indices: Final = self.model_name_to_deployment_indices.get(model_group_name)
# Return first deployment for this model_name
model: Final = self.model_list[indices[0]] if indices else None
if model is not None:
if isinstance(model, dict):
return Deployment(**model)
elif isinstance(model, Deployment):
return model
else:
raise Exception(f"Model Name invalid - {type(model)}")
return None
@staticmethod
@ -11759,11 +11804,12 @@ class Router:
by other teams so a shared model name never resolves another team's
credentials.
"""
indices: Final = self.model_name_to_deployment_indices.get(model_group_name) or ()
usable: Final = (
self.model_list[idx] for idx in indices if self.deployment_usable_by_team(self.model_list[idx], team_id)
)
first_usable: Final = next(usable, None)
with self._deployment_index_lock:
indices: Final = self.model_name_to_deployment_indices.get(model_group_name) or ()
usable: Final = (
self.model_list[idx] for idx in indices if self.deployment_usable_by_team(self.model_list[idx], team_id)
)
first_usable: Final = next(usable, None)
if first_usable is None:
return None
return Deployment(**first_usable) if isinstance(first_usable, dict) else first_usable
@ -11866,11 +11912,10 @@ class Router:
this never triggers pattern matching or deep copies, so it is safe to call per
listed model on the /v1/models hot path.
"""
indices: Final = self.model_name_to_deployment_indices.get(model_name)
if not indices:
deployments: Final = self.deployments_named(model_name)
if not deployments:
return None
deployments: Final = tuple(self.model_list[index] for index in indices)
model_infos: Final = tuple(
MappingProxyType(
{
@ -12043,10 +12088,11 @@ class Router:
def _get_team_public_name_deployment(self, model_id: str, team_id: str | None) -> Deployment | None:
if team_id is None:
return None
team_indices: Final = self.team_model_to_deployment_indices.get((team_id, model_id))
if not team_indices:
with self._deployment_index_lock:
team_indices: Final = self.team_model_to_deployment_indices.get((team_id, model_id))
team_model: Final = self.model_list[team_indices[0]] if team_indices else None
if team_model is None:
return None
team_model: Final = self.model_list[team_indices[0]]
return Deployment(**team_model) if isinstance(team_model, dict) else team_model
def _get_wildcard_deployment_usable_by_team(self, model_id: str, team_id: str | None) -> Deployment | None:
@ -12285,10 +12331,9 @@ class Router:
Optimized with O(1) index lookup instead of O(n) linear scan.
"""
# O(1) lookup via model_id_to_deployment_index_map
if id in self.model_id_to_deployment_index_map:
idx: Final = self.model_id_to_deployment_index_map[id]
return self.model_list[idx]
return None
with self._deployment_index_lock:
idx: Final = self.model_id_to_deployment_index_map.get(id)
return None if idx is None else self.model_list[idx]
def get_model_group(self, id: str) -> list | None:
"""
@ -12381,6 +12426,17 @@ class Router:
)
return model_info
@staticmethod
def _model_group_llm_provider(litellm_params: LiteLLM_Params) -> tuple[str, str]:
try:
return resolve_model_and_provider_for_metadata(
model=litellm_params.model,
custom_llm_provider=litellm_params.custom_llm_provider,
)
except litellm.exceptions.BadRequestError as e:
verbose_router_logger.error("litellm.router.py::get_model_group_info() - %s", e)
return "", ""
def _set_model_group_info(
self, model_group: str, user_facing_model_group_name: str, team_id: str | None = None
) -> ModelGroupInfo | None:
@ -12469,15 +12525,7 @@ class Router:
deployment_is_mapped = deployment_is_catalog_mapped(model_info, model_info_dict)
# get llm provider
litellm_model, llm_provider = "", ""
try:
litellm_model, llm_provider, _, _ = litellm.get_llm_provider(
model=litellm_params.model,
custom_llm_provider=litellm_params.custom_llm_provider,
litellm_params=litellm_params,
)
except litellm.exceptions.BadRequestError as e:
verbose_router_logger.error("litellm.router.py::get_model_group_info() - %s", e)
litellm_model, llm_provider = self._model_group_llm_provider(litellm_params)
if model_info is None:
supported_openai_params = litellm.get_supported_openai_params(
@ -12958,20 +13006,19 @@ class Router:
if model_name is not None:
# O(1) lookup in model_name index, then O(k) iteration where k = deployments for this model_name
if model_name in self.model_name_to_deployment_indices:
indices: Final = self.model_name_to_deployment_indices[model_name]
for idx in indices:
model = self.model_list[idx]
if "model_info" in model and "id" in model["model_info"]:
if exclude_team_models and model["model_info"].get("team_id"):
continue
ids.append(model["model_info"]["id"])
for model in self.deployments_named(model_name) or ():
if "model_info" in model and "id" in model["model_info"]:
if exclude_team_models and model["model_info"].get("team_id"):
continue
ids.append(model["model_info"]["id"])
else:
# When model_name is None, return all model IDs
# Use the index map keys for O(n) where n = total deployments
for model_id in self.model_id_to_deployment_index_map:
idx = self.model_id_to_deployment_index_map[model_id]
model = self.model_list[idx]
with self._deployment_index_lock:
id_indices: Final = tuple(self.model_id_to_deployment_index_map.items())
models: Final = tuple(self.model_list)
for model_id, idx in id_indices:
model = models[idx]
if "model_info" in model and "id" in model["model_info"]:
if exclude_team_models and model["model_info"].get("team_id"):
continue
@ -13249,7 +13296,9 @@ class Router:
# handled downstream by the pattern_router in _common_checks_available_deployment.
return None
def should_include_deployment(self, model_name: str, model: dict, team_id: str | None = None) -> bool:
def should_include_deployment(
self, model_name: str, model: DeploymentTypedDict | dict, team_id: str | None = None
) -> bool:
"""
Get the team-specific model name if team_id matches the deployment.
"""
@ -13297,12 +13346,10 @@ class Router:
# O(1) lookup in team_model index when team_id is provided
if team_id is not None:
key: Final = (team_id, model_name)
if key in self.team_model_to_deployment_indices:
indices = self.team_model_to_deployment_indices[key]
team_models: Final = self._team_deployments_named(team_id, model_name)
if team_models is not None:
# O(k) where k = team deployments for this model_name (typically 1-10)
for idx in indices:
model = self.model_list[idx]
for model in team_models:
if not self.should_include_deployment(model_name=model_name, model=model, team_id=team_id):
continue
if model_alias is not None:
@ -13315,12 +13362,10 @@ class Router:
return returned_models
# O(1) lookup in model_name index
if model_name in self.model_name_to_deployment_indices:
indices = self.model_name_to_deployment_indices[model_name]
named_models: Final = self.deployments_named(model_name)
if named_models is not None:
# O(k) where k = deployments for this model_name (typically 1-10)
for idx in indices:
model = self.model_list[idx]
for model in named_models:
if self.should_include_deployment(model_name=model_name, model=model, team_id=team_id):
if model_alias is not None:
# Optimized: Use shallow copy since we only modify top-level model_name
@ -13334,7 +13379,7 @@ class Router:
# Fallback: if team_id is provided and model_name not in index,
# check if model_name matches any team_public_model_name
# O(n) scan but only when team_id lookup fails
for idx, model in enumerate(self.model_list):
for _idx, model in enumerate(self.model_list):
if self.should_include_deployment(model_name=model_name, model=model, team_id=team_id):
if model_alias is not None:
# Optimized: Use shallow copy since we only modify top-level model_name
@ -14599,12 +14644,13 @@ class Router:
def _team_deployments_across_teams(self, model: str) -> list[DeploymentTypedDict]:
"""Every team's deployments under public name `model`, for a proxy admin calling without a team."""
team_deployments: Final = [
self.model_list[index]
for (_, public_model_name), indices in self.team_model_to_deployment_indices.items()
if public_model_name == model
for index in indices
]
with self._deployment_index_lock:
team_deployments: Final = [
self.model_list[index]
for (_, public_model_name), indices in self.team_model_to_deployment_indices.items()
if public_model_name == model
for index in indices
]
team_ids: Final = {
team_id
for deployment in team_deployments

View file

@ -3126,17 +3126,14 @@ class ComplexityRouter(CustomLogger):
model_to_prefs: Final[dict[str, AdaptiveRouterPreferences]] = {}
model_to_cost: Final[dict[str, float]] = {}
model_list: Final = getattr(self.litellm_router_instance, "model_list", None) or []
name_to_indices: Final = getattr(self.litellm_router_instance, "model_name_to_deployment_indices", {}) or {}
for name in available_models:
indices = name_to_indices.get(name, [])
if not indices:
deployments = self.litellm_router_instance.deployments_named(name)
if not deployments:
model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[])
model_to_cost[name] = 0.0
continue
deployment = model_list[indices[0]]
mi = deployment.get("model_info") if isinstance(deployment, dict) else deployment.model_info
mi_dict: dict[str, Any] = mi if isinstance(mi, dict) else (mi.model_dump() if mi else {})
deployment = deployments[0]
mi_dict: dict[str, Any] = deployment.get("model_info") or {}
prefs_raw = mi_dict.get("adaptive_router_preferences")
if prefs_raw is not None:
model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw)
@ -3144,9 +3141,7 @@ class ComplexityRouter(CustomLogger):
model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[])
# model_info is the conventional pricing location elsewhere in LiteLLM; litellm_params wins if set.
lp = deployment.get("litellm_params") if isinstance(deployment, dict) else deployment.litellm_params
lp_dict: dict[str, Any] = lp if isinstance(lp, dict) else (lp.model_dump() if lp else {})
cost = lp_dict.get("input_cost_per_token")
cost = deployment.get("litellm_params", {}).get("input_cost_per_token")
if cost is None:
cost = mi_dict.get("input_cost_per_token")
model_to_cost[name] = float(cost) if cost is not None else 0.0

View file

@ -14,7 +14,7 @@ bare string with no provider beside them; an operator who writes ``deepseek-r1``
Azure would otherwise be priced against whoever else owns that name.
"""
from collections.abc import Iterable
from collections.abc import Iterable, Mapping
from typing import TYPE_CHECKING, Final, NamedTuple
from litellm._logging import verbose_router_logger
@ -80,12 +80,11 @@ def _models_in(router: "Router", group_name: str) -> tuple[Baseline, ...]:
straight at a provider model rather than at a configured group, and prices under
its own name because there is no deployment to override it.
"""
indices: Final = router.model_name_to_deployment_indices.get(group_name)
if not indices:
deployments: Final = router.deployments_named(group_name)
if not deployments:
return (Baseline(qualified),) if (qualified := canonical_model(group_name)) else ()
def candidate(index: int) -> Baseline | None:
deployment: Final = router.model_list[index]
def candidate(deployment: Mapping[str, object]) -> Baseline | None:
params: Final = deployment.get("litellm_params")
if not isinstance(params, dict):
return None
@ -98,7 +97,7 @@ def _models_in(router: "Router", group_name: str) -> tuple[Baseline, ...]:
deployment_id: Final = info.get("id") if isinstance(info, dict) else None
return Baseline(qualified, str(deployment_id) if deployment_id else None)
return tuple(c for index in indices if (c := candidate(index)) is not None)
return tuple(c for deployment in deployments if (c := candidate(deployment)) is not None)
def _priced(router: "Router", candidate: Baseline) -> tuple[float, Baseline] | None:

View file

@ -91,6 +91,7 @@ class ProviderConnection:
litellm_credential_name: str | None = None
configurable_clientside_auth_params: "Sequence[str | ConfigurableClientsideParamsCustomAuth] | None" = None
use_xai_oauth: bool | None = None
chatgpt_auth_file: str | None = None
github_copilot_auth_type: str | None = None
github_copilot_host: str | None = None
github_copilot_user_session: object | None = None

View file

@ -55,6 +55,8 @@ class LiteLLMKeyType(str, enum.Enum):
class GenerateKeyRequest(KeyRequestBase):
tpm_limit: int | None = Field(default=None, ge=0)
rpm_limit: int | None = Field(default=None, ge=0)
soft_budget: float | None = None
send_invite_email: bool | None = None
key_type: LiteLLMKeyType | None = Field(
@ -82,6 +84,8 @@ def _is_mapping(value: object) -> TypeIs[Mapping[object, object]]: # guard-ok:
class UpdateKeyRequest(KeyRequestBase):
tpm_limit: int | None = Field(default=None, ge=0)
rpm_limit: int | None = Field(default=None, ge=0)
# Note: the defaults of all Params here MUST BE NONE
# else they will get overwritten
duration: str | None = None

View file

@ -16,6 +16,8 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase):
class NewUserRequest(GenerateRequestBase):
tpm_limit: int | None = Field(default=None, ge=0)
rpm_limit: int | None = Field(default=None, ge=0)
max_budget: float | None = None
user_email: str | None = None
user_alias: str | None = None
@ -47,6 +49,8 @@ class NewUserRequest(GenerateRequestBase):
class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with BulkUpdateUserRequest
tpm_limit: int | None = Field(default=None, ge=0)
rpm_limit: int | None = Field(default=None, ge=0)
# repr=False keeps the plaintext out of management-endpoint alerts, which str() the request model
password: str | None = Field(default=None, repr=False)
spend: float | None = None

View file

@ -572,6 +572,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
default=False,
description="Use stored xAI OAuth credentials when no xAI API key is configured.",
)
chatgpt_auth_file: str | None = Field(
default=None,
description="Path to a ChatGPT OAuth auth.json for this deployment.",
)
fireworks_forward_user_id: bool | None = Field(
default=None,
description="Send the LiteLLM user id of the calling key as the `user` field on Fireworks AI chat, responses and messages requests.",

View file

@ -4360,6 +4360,7 @@ server_owned_wif_litellm_params: Final = (
+ openai_wif_litellm_params
+ oauth_token_exchange_litellm_params
+ github_copilot_oauth_litellm_params
+ ("chatgpt_auth_file",)
)
secret_bearing_wif_litellm_params: Final = tuple(sorted(WIF_SECRET_BEARING_KEYS))

View file

@ -2964,12 +2964,14 @@ def supports_native_streaming(model: str, custom_llm_provider: str | None) -> bo
Raises:
Exception: If the given model is not found in model_prices_and_context_window.json.
"""
from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata
try:
model, custom_llm_provider, _, _ = litellm.get_llm_provider(
resolved_model, resolved_provider = resolve_model_and_provider_for_metadata(
model=model, custom_llm_provider=custom_llm_provider
)
model_info: Final = get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)
model_info: Final = get_model_info_helper(model=resolved_model, custom_llm_provider=resolved_provider)
supports_native_streaming = model_info.get("supports_native_streaming", True)
if supports_native_streaming is None:
supports_native_streaming = True

View file

@ -6206,6 +6206,8 @@
"text",
"audio"
],
"supports_audio_input": true,
"supports_audio_output": true,
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
@ -6270,6 +6272,8 @@
"text",
"audio"
],
"supports_audio_input": true,
"supports_audio_output": true,
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
@ -33698,7 +33702,7 @@
"input_cost_per_token": 1e-05,
"input_cost_per_token_above_272k_tokens": 2e-05,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -33722,7 +33726,7 @@
"input_cost_per_token": 1e-07,
"input_cost_per_token_above_272k_tokens": 2e-07,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -33746,7 +33750,7 @@
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_272k_tokens": 4e-06,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -33770,7 +33774,7 @@
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_272k_tokens": 4e-06,
"litellm_provider": "github_copilot",
"max_input_tokens": 272000,
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
@ -78851,5 +78855,47 @@
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true
},
"baseten/decision-models/baseten-elephant-9b-latest": {
"input_cost_per_token": 0.0,
"litellm_provider": "baseten",
"max_input_tokens": 16384,
"mode": "decisions",
"output_cost_per_token": 0.0,
"source": "https://inference.baseten.co/v1/models",
"supports_vision": true
},
"baseten/decision-models/baseten-mammoth-27b-latest": {
"input_cost_per_token": 0.0,
"litellm_provider": "baseten",
"max_input_tokens": 16384,
"mode": "decisions",
"output_cost_per_token": 0.0,
"source": "https://inference.baseten.co/v1/models",
"supports_vision": true
},
"github_copilot/gpt-6.1-terra": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-06,
"cache_read_input_token_cost": 5e-08,
"cache_read_input_token_cost_above_272k_tokens": 1e-07,
"input_cost_per_token": 1e-06,
"input_cost_per_token_above_272k_tokens": 2e-06,
"litellm_provider": "github_copilot",
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 5e-06,
"output_cost_per_token_above_272k_tokens": 7.5e-06,
"source": "https://github.blog/changelog/2026-10-26-gpt-6-1-terra-in-github-copilot",
"supported_endpoints": [
"/v1/responses"
],
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_vision": true
}
}

View file

@ -377,21 +377,8 @@ required-version = ">=0.10.9"
exclude-newer = "3 days"
[tool.uv.exclude-newer-package]
# TODO: ryan, remove after 2026-10-23, once certifi 2026.10.19 (drops the compromised eMudhra root) clears the 3-day window
certifi = "2026-10-20T03:00:00Z"
# TODO(shrey): remove after 2026-10-25, once grpcio 1.84.1 (drops the compromised eMudhra roots from roots.pem) clears the 3-day window
grpcio = "2026-10-22T15:00:00Z"
# TODO(shrey): remove after 2026-10-24, once boto3/botocore 1.43.117 (drops the compromised eMudhra roots from cacert.pem) clear the 3-day window
boto3 = "2026-10-21T19:31:00Z"
botocore = "2026-10-21T19:31:00Z"
# TODO: remove after 2026-10-25, once brotli 1.2.1 (CVE-2026-111497) clears the 3-day window
brotli = "2026-10-22T14:01:00Z"
# TODO: ryan, remove after 2026-10-24, once mlflow 3.17.2 (GHSA-wqmh-qccw-pqwf) clears the 3-day window
mlflow = "2026-10-21T03:00:00Z"
mlflow-skinny = "2026-10-21T03:00:00Z"
mlflow-tracing = "2026-10-21T03:00:00Z"
# TODO: ryan, remove after 2026-10-25, once pypdf 6.21.2 (GHSA-6c7f-9265-g8fw) clears the 3-day window
pypdf = "2026-10-21T17:00:00Z"
[tool.uv.sources]
litellm-proxy-extras = { workspace = true }

View file

@ -1,547 +0,0 @@
"""Sync the together_ai entries of model_prices_and_context_window.json with Together's live serverless catalog.
Pulls ``GET https://api.together.ai/v1/models?serverless`` plus the deprecations doc, maps API fields onto
registry fields, merges the reviewed capability rules below for everything the API cannot express, and diffs
the result against the registry. Dry run (the default) prints the diff summary and the generated PR body;
``--write`` applies the changes to the root cost map and its ``litellm/`` backup copy.
Policy highlights:
- Prices arrive per 1M tokens with float artifacts and are normalized to clean per-token values.
- A registry entry absent from the serverless catalog is marked with ``deprecation_date`` from the docs
deprecation table, never deleted; absences with no docs date are surfaced for a human call.
- Availability comes from the API: a model the docs list as removed but the API still serves stays live,
with the conflict surfaced as a warning.
- Manually curated values the API cannot express (``metadata.successor``, ``max_output_tokens`` on existing
entries, capability flags no rule covers) are never overwritten; conflicts are surfaced instead.
"""
import argparse
import json
import os
import re
import sys
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from pathlib import Path
from types import MappingProxyType
from typing import Final
import httpx
from pydantic import BaseModel, TypeAdapter, ValidationError
MODELS_URL: Final = "https://api.together.ai/v1/models?serverless"
DEPRECATIONS_URL: Final = "https://docs.together.ai/docs/deprecations.md"
PROVIDER: Final = "together_ai"
PREFIX: Final = "together_ai/"
SOURCE_URL: Final = "https://docs.together.ai/docs/serverless-models"
COST_MAP_RELPATHS: Final = (
"model_prices_and_context_window.json",
"litellm/model_prices_and_context_window_backup.json",
)
TYPE_TO_MODE: Final = MappingProxyType({"chat": "chat", "embedding": "embedding", "moderation": "chat"})
class SyncError(RuntimeError):
pass
class CatalogPricing(BaseModel):
input: float
output: float
cached_input: float | None = None
class CatalogModel(BaseModel):
id: str
type: str
context_length: int | None = None
pricing: CatalogPricing
CATALOG_ADAPTER: Final = TypeAdapter(list[CatalogModel])
RegistryEntry = dict[str, object]
CostMap = dict[str, object]
@dataclass(frozen=True, slots=True)
class CapabilityRule:
model_id: str
fields: Mapping[str, bool | int]
provenance: str
def _rule(model_id: str, provenance: str, **fields: bool | int) -> CapabilityRule:
return CapabilityRule(model_id=model_id, fields=MappingProxyType(dict(fields)), provenance=provenance)
_TOOLS: Final = MappingProxyType(
{
"supports_function_calling": True,
"supports_parallel_function_calling": True,
"supports_response_schema": True,
"supports_tool_choice": True,
}
)
CAPABILITY_RULES: Final = (
_rule(
"MiniMaxAI/MiniMax-M3",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/minimax-m3",
**_TOOLS,
supports_reasoning=True,
supports_vision=True,
),
_rule("Prism-ML/Ternary-Bonsai-27B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule(
"Qwen/Qwen3.5-9B",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/qwen3-5-9b",
**_TOOLS,
supports_reasoning=True,
supports_vision=True,
),
_rule(
"Qwen/Qwen3.6-Plus",
"reviewed for the LIT-5968 backfill; hybrid reasoning model without a documented tools contract",
supports_reasoning=True,
),
_rule("Qwen/Qwen3.7-Max", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule("Qwen/Qwen3.7-Plus", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule("Qwen/Qwen3.8-2.4T-A95B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule("arize-ai/qwen-2-1.5b-instruct", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule(
"deepseek-ai/DeepSeek-V4-Flash-0731",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-flash",
**_TOOLS,
),
_rule(
"deepseek-ai/DeepSeek-V4-Pro",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-pro",
**_TOOLS,
supports_reasoning=True,
),
_rule(
"deepseek-ai/DeepSeek-V4-Pro-0813",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-pro",
**_TOOLS,
),
_rule("google/gemma-3n-E4B-it", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule(
"google/gemma-4-31B-it",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/gemma-4-31b-it",
**_TOOLS,
supports_vision=True,
),
_rule(
"intfloat/multilingual-e5-large-instruct",
"embedding dims per https://huggingface.co/intfloat/multilingual-e5-large-instruct",
output_vector_size=1024,
),
_rule(
"meta-llama/Llama-3.3-70B-Instruct-Turbo",
"reviewed for the LIT-5968 backfill against https://docs.together.ai/docs/function-calling",
**_TOOLS,
),
_rule(
"meta-llama/Llama-Guard-4-12B",
"moderation classifier with a chat-shaped API; no tools per the LIT-5968 backfill review",
),
_rule("meta-models/Muse-Glimmer-30B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule(
"moonshotai/Kimi-K2.7-Code",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/kimi-k2-7-code",
**_TOOLS,
supports_vision=True,
),
_rule(
"moonshotai/Kimi-K3",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/kimi-k3",
**_TOOLS,
supports_reasoning=True,
supports_vision=True,
),
_rule(
"nvidia/nemotron-3-ultra-550b-a55b",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/nemotron-3-ultra",
**_TOOLS,
supports_reasoning=True,
),
_rule(
"openai/gpt-oss-120b",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/gpt-oss-120b",
**_TOOLS,
supports_reasoning=True,
),
_rule(
"openai/gpt-oss-20b",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/gpt-oss-20b",
**_TOOLS,
),
_rule("pearl-ai/gemma-4-31b-it", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule(
"thinkingmachines/Inkling",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/inkling",
**_TOOLS,
),
_rule("thinkingmachines/Inkling-Small", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
_rule(
"zai-org/GLM-5.2",
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/glm-5-2;"
" 128K output ceiling per https://docs.z.ai/guides/llm/glm-5.2",
**_TOOLS,
supports_reasoning=True,
max_output_tokens=128000,
max_tokens=128000,
),
_rule(
"zai-org/GLM-5.3-Flash",
"reviewed for LIT-6489 against https://www.together.ai/models/glm-5-3-flash;"
" 128K output ceiling per https://docs.z.ai/guides/llm/glm-5.3",
**_TOOLS,
supports_reasoning=True,
supports_vision=True,
max_output_tokens=128000,
max_tokens=128000,
),
)
RULES_BY_ID: Final = MappingProxyType({rule.model_id: rule for rule in CAPABILITY_RULES})
@dataclass(frozen=True, slots=True)
class DeprecationDoc:
removal_dates: Mapping[str, str]
redirects: Mapping[str, str]
_REDIRECT_ROW: Final = re.compile(r"^\|\s*`([^`]+)`\s*\|\s*`([^`]+)`\s*\|")
_REMOVAL_ROW: Final = re.compile(r"^\|\s*(\d{4}-\d{2}-\d{2})\s*\|\s*`([^`]+)`\s*\|")
def _section(markdown: str, heading: str) -> str:
level: Final = heading.split(" ", 1)[0]
start: Final = markdown.find(f"\n{heading}\n")
if start < 0:
return ""
body: Final = markdown[start + 1 + len(heading) :]
next_heading: Final = re.search(rf"^{re.escape(level)} ", body, flags=re.MULTILINE)
return body[: next_heading.start()] if next_heading else body
def parse_deprecations(markdown: str) -> DeprecationDoc:
redirect_rows: Final = tuple(
m.groups()
for m in (_REDIRECT_ROW.match(line) for line in _section(markdown, "## Active model redirects").splitlines())
if m
)
inference: Final = _section(_section(markdown, "## Deprecation history"), "### Inference")
removal_rows: Final = tuple(m.groups() for m in (_REMOVAL_ROW.match(line) for line in inference.splitlines()) if m)
if not redirect_rows or not removal_rows:
raise SyncError(
"deprecations doc parsed to zero redirect or removal rows; the table format at "
f"{DEPRECATIONS_URL} changed and the parser needs updating"
)
removal_dates: Final = {model: date for date, model in reversed(removal_rows)}
return DeprecationDoc(
removal_dates=MappingProxyType(dict(reversed(removal_dates.items()))),
redirects=MappingProxyType({original: target for original, target in redirect_rows}),
)
def per_token(price_per_million: float) -> float:
return float(f"{price_per_million / 1e6:.6g}")
def _resolve_name(name: str, universe: frozenset[str]) -> str | None:
if name in universe:
return name
suffix_matches: Final = tuple(candidate for candidate in universe if candidate.endswith(f"/{name}"))
return suffix_matches[0] if len(suffix_matches) == 1 else None
def resolve_successor(model_id: str, doc: DeprecationDoc, live_ids: frozenset[str]) -> str | None:
canonical: Final = live_ids | frozenset(doc.removal_dates)
redirects: Final = {
(_resolve_name(raw_source, canonical) or raw_source): (_resolve_name(raw_target, canonical) or raw_target)
for raw_source, raw_target in doc.redirects.items()
}
seen: Final = set()
current = model_id # rebind-ok: walks the redirect chain
while current in redirects and current not in seen:
seen.add(current)
current = redirects[current] # rebind-ok: walks the redirect chain
return current if current != model_id and current in live_ids else None
@dataclass(frozen=True, slots=True)
class SyncOutcome:
cost_map: CostMap
added: tuple[str, ...] = ()
updated: tuple[str, ...] = ()
deprecated: tuple[str, ...] = ()
reappeared: tuple[str, ...] = ()
warnings: tuple[str, ...] = ()
skipped_types: Mapping[str, int] = field(default_factory=dict)
@property
def has_changes(self) -> bool:
return bool(self.added or self.updated or self.deprecated or self.reappeared)
def _api_fields(model: CatalogModel) -> RegistryEntry:
cached: Final = model.pricing.cached_input
return {
"input_cost_per_token": per_token(model.pricing.input),
"output_cost_per_token": per_token(model.pricing.output),
**({"cache_read_input_token_cost": per_token(cached), "supports_prompt_caching": True} if cached else {}),
**({"max_input_tokens": model.context_length} if model.context_length is not None else {}),
}
def _new_entry(model: CatalogModel, mode: str) -> RegistryEntry:
rule: Final = RULES_BY_ID.get(model.id)
legacy_ceiling: Final = {} if model.context_length is None else {"max_tokens": model.context_length}
merged: Final = {
**_api_fields(model),
**legacy_ceiling,
"litellm_provider": PROVIDER,
"mode": mode,
"source": SOURCE_URL,
**(dict(rule.fields) if rule else {}),
}
return dict(sorted(merged.items()))
def _updated_entry(entry: RegistryEntry, model: CatalogModel) -> tuple[RegistryEntry, tuple[str, ...]]:
rule: Final = RULES_BY_ID.get(model.id)
desired: Final = {**_api_fields(model), **(dict(rule.fields) if rule else {})}
dropped: Final = () if model.pricing.cached_input else ("cache_read_input_token_cost", "supports_prompt_caching")
changes: Final = tuple(
f"{name}: {entry.get(name)!r} -> {value!r}" for name, value in desired.items() if entry.get(name) != value
) + tuple(
f"{name}: {entry[name]!r} removed (no longer in the catalog pricing)" for name in dropped if name in entry
)
merged: Final = {name: value for name, value in {**entry, **desired}.items() if name not in dropped}
return dict(sorted(merged.items())), changes
def _with_new_keys_in_block(original: CostMap, result: CostMap, new_keys: Sequence[str]) -> CostMap:
provider_keys: Final = tuple(key for key in original if key.startswith(PREFIX))
if not new_keys or not provider_keys:
return result
block_end: Final = provider_keys[-1]
return {
key: value
for existing in original
for key, value in (
(existing, result[existing]),
*((new, result[new]) for new in sorted(new_keys) if existing == block_end),
)
}
def compute_sync(cost_map: CostMap, catalog: Sequence[CatalogModel], doc: DeprecationDoc) -> SyncOutcome:
live_ids: Final = frozenset(model.id for model in catalog)
token_models: Final = {model.id: model for model in catalog if model.type in TYPE_TO_MODE}
skipped: Final = {
model.type: sum(1 for m in catalog if m.type == model.type)
for model in catalog
if model.type not in TYPE_TO_MODE
}
registry_ids: Final = {key.removeprefix(PREFIX): key for key in cost_map if key.startswith(PREFIX)}
added: Final[list[str]] = []
updated: Final[list[str]] = []
deprecated: Final[list[str]] = []
reappeared: Final[list[str]] = []
warnings: Final[list[str]] = []
result: Final[CostMap] = dict(cost_map)
for model_id, model in sorted(token_models.items()):
mode: Final = TYPE_TO_MODE[model.type]
key: Final = f"{PREFIX}{model_id}"
if model_id in doc.removal_dates:
warnings.append(
f"`{key}` is listed as removed on {doc.removal_dates[model_id]} in the docs but the serverless "
"catalog still serves it; availability kept from the API"
)
entry = result.get(key)
if not isinstance(entry, dict):
result[key] = _new_entry(model, mode)
added.append(key)
if model.type == "chat" and model_id not in RULES_BY_ID:
warnings.append(
f"`{key}` added without a capability rule; review its tools/vision/reasoning support and add one"
)
continue
if entry.get("mode") != mode:
warnings.append(
f"`{key}` has curated mode {entry.get('mode')!r} but the catalog maps to {mode!r}; left unchanged"
)
new_entry, changes = _updated_entry(entry, model)
if "deprecation_date" in new_entry:
new_entry.pop("deprecation_date")
reappeared.append(key)
if changes:
updated.append(f"{key}: " + "; ".join(changes))
if changes or key in reappeared:
result[key] = new_entry
for model_id, key in sorted(registry_ids.items()):
if model_id in token_models:
continue
entry = result.get(key)
if not isinstance(entry, dict):
continue
removal_date: Final = doc.removal_dates.get(model_id)
successor: Final = resolve_successor(model_id, doc, live_ids)
metadata = entry.get("metadata")
curated_successor: Final = metadata.get("successor") if isinstance(metadata, dict) else None
new_entry = dict(entry)
if removal_date is not None and entry.get("deprecation_date") != removal_date:
if "deprecation_date" in entry:
warnings.append(
f"`{key}` has curated deprecation_date {entry.get('deprecation_date')!r} but the docs list "
f"{removal_date!r}; left unchanged"
)
else:
new_entry["deprecation_date"] = removal_date
if removal_date is None and "deprecation_date" not in entry:
warnings.append(
f"`{key}` is absent from the serverless catalog with no removal date in the docs; "
"needs a human deprecation call"
)
if successor is not None:
desired_successor: Final = f"{PREFIX}{successor}"
if curated_successor is None:
new_entry["metadata"] = dict(
sorted({**(metadata if isinstance(metadata, dict) else {}), "successor": desired_successor}.items())
)
elif curated_successor != desired_successor:
warnings.append(
f"`{key}` has curated successor {curated_successor!r} but the docs redirects resolve to "
f"{desired_successor!r}; left unchanged"
)
if new_entry != entry:
result[key] = dict(sorted(new_entry.items()))
deprecated.append(f"{key}: " + ", ".join(sorted(set(new_entry) - set(entry)) or ["updated"]))
return SyncOutcome(
cost_map=_with_new_keys_in_block(cost_map, result, tuple(added)),
added=tuple(added),
updated=tuple(updated),
deprecated=tuple(deprecated),
reappeared=tuple(reappeared),
warnings=tuple(warnings),
skipped_types=MappingProxyType(skipped),
)
def _section_block(title: str, lines: Sequence[str], backtick: bool) -> str:
bullets: Final = "\n".join(f"- `{line}`" if backtick else f"- {line}" for line in lines) or "- none"
return f"### {title} ({len(lines)})\n{bullets}\n"
def render_pr_body(outcome: SyncOutcome) -> str:
skipped: Final = ", ".join(f"{kind} ({count})" for kind, count in sorted(outcome.skipped_types.items())) or "none"
return (
"Automated daily sync of the together_ai entries in model_prices_and_context_window.json against "
f"`GET {MODELS_URL}` and {DEPRECATIONS_URL} by scripts/sync_together_ai_models.py.\n"
"\n"
f"{_section_block('Added', outcome.added, backtick=True)}"
"\n"
f"{_section_block('Updated', outcome.updated, backtick=True)}"
"\n"
f"{_section_block('Marked deprecated', outcome.deprecated, backtick=True)}"
"\n"
f"{_section_block('Returned to the catalog', outcome.reappeared, backtick=True)}"
"\n"
f"{_section_block('Warnings needing a human call', outcome.warnings, backtick=False)}"
"\n"
f"Catalog model types outside the sync's token-pricing scope, skipped: {skipped}\n"
)
def render_summary(outcome: SyncOutcome) -> str:
return (
f"added={len(outcome.added)} updated={len(outcome.updated)} deprecated={len(outcome.deprecated)} "
f"reappeared={len(outcome.reappeared)} warnings={len(outcome.warnings)}"
)
def load_catalog(raw: bytes) -> list[CatalogModel]:
parsed: Final = json.loads(raw)
entries: Final = parsed.get("data") if isinstance(parsed, dict) else parsed
try:
catalog: Final = CATALOG_ADAPTER.validate_python(entries)
except ValidationError as error:
raise SyncError(f"the catalog response no longer matches the expected shape: {error}") from error
if not any(model.type in TYPE_TO_MODE for model in catalog):
raise SyncError(
"the catalog response contains no token-priced models; refusing to mark the whole registry deprecated"
)
return catalog
def _fetch(url: str, headers: Mapping[str, str]) -> bytes:
response: Final = httpx.get(url, headers=dict(headers), timeout=30, follow_redirects=True)
if response.status_code != 200:
raise SyncError(f"GET {url} returned {response.status_code}")
return response.content
def _serialize(cost_map: CostMap) -> str:
return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n"
def main(argv: Sequence[str]) -> int:
parser: Final = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--write", action="store_true", help="apply the sync to the cost map files (default: dry run)")
parser.add_argument("--models-json", type=Path, help="recorded catalog response to use instead of the live API")
parser.add_argument(
"--deprecations-md", type=Path, help="recorded deprecations doc to use instead of the live docs"
)
parser.add_argument("--pr-body-file", type=Path, help="write the generated PR body to this path")
parser.add_argument("--repo-root", type=Path, default=Path(__file__).resolve().parent.parent)
args: Final = parser.parse_args(argv)
if args.models_json is not None:
catalog_raw: Final = args.models_json.read_bytes()
else:
api_key: Final = os.environ.get("TOGETHER_API_KEY")
if not api_key:
raise SyncError("TOGETHER_API_KEY is not set and --models-json was not given")
catalog_raw = _fetch(MODELS_URL, {"Authorization": f"Bearer {api_key}"}) # rebind-ok: branch-dependent source
catalog: Final = load_catalog(catalog_raw)
markdown: Final = (
args.deprecations_md.read_text() if args.deprecations_md is not None else _fetch(DEPRECATIONS_URL, {}).decode()
)
doc: Final = parse_deprecations(markdown)
cost_map_path: Final = args.repo_root / COST_MAP_RELPATHS[0]
cost_map: Final = json.loads(cost_map_path.read_text())
outcome: Final = compute_sync(cost_map, catalog, doc)
body: Final = render_pr_body(outcome)
if args.pr_body_file is not None:
args.pr_body_file.write_text(body)
if args.write and outcome.has_changes:
for relpath in COST_MAP_RELPATHS:
(args.repo_root / relpath).write_text(_serialize(outcome.cost_map))
print(render_summary(outcome))
print()
print(body)
if not args.write:
print("dry run: no files were touched")
elif not outcome.has_changes:
print("registry already in sync: no files were touched")
return 0
if __name__ == "__main__":
try:
raise SystemExit(main(sys.argv[1:]))
except SyncError as error:
print(f"SYNC FAILED: {error}", file=sys.stderr)
raise SystemExit(1) from error

View file

@ -171,6 +171,7 @@ ignored_function_names = [
"_deployment_field", # Read by is_batch_capable, tested through the mixed-group tests in test_router.py
"_deployment_provider", # Read by without_deployments_that_cannot_batch, tested through the mixed-group tests in test_router.py
"_cannot_run_batches", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py
"_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py
]

View file

@ -2,7 +2,7 @@ import { chromium, expect, request } from "@playwright/test";
import { users, Role, STORAGE_PATHS } from "./fixtures/users";
import { ARTIFACT_DIR, UI_BASE_URL } from "./constants";
import { expectUnrestrictedDashboard, setInvitedUserPassword } from "./helpers/userOnboarding";
import { hideLiteAdmin } from "./helpers/navigation";
import { setAriaHidden } from "./helpers/navigation";
import { masterKey as getMasterKey } from "./helpers/traffic";
import * as fs from "fs";
import * as path from "path";
@ -78,7 +78,7 @@ async function globalSetup() {
await dismiss.click();
}
if (role === Role.ProxyAdmin) {
await hideLiteAdmin(page);
await setAriaHidden(page, true);
}
// The login flow stores a post-login return URL in the litellm_return_url
// cookie. If the snapshot captures it before the app consumes it, every

View file

@ -62,18 +62,29 @@ export async function dismissFeedbackPopup(page: PlaywrightPage): Promise<void>
}
}
export async function hideLiteAdmin(page: PlaywrightPage): Promise<void> {
export async function setAriaHidden(page: PlaywrightPage, hidden: boolean): Promise<void> {
await page.getByRole("button", { name: /Account menu/i }).click();
const panel = page.getByTestId("sidebar-account-menu-panel");
await expect(panel).toBeVisible({ timeout: 5_000 });
const toggle = panel.getByRole("switch", { name: "Toggle hide LiteAdmin" });
if ((await toggle.getAttribute("aria-checked")) !== "true") {
await toggle.click();
const premium = await panel.getByText("Premium", { exact: true }).isVisible();
const toggle = panel.getByRole("switch", { name: "Toggle hide Aria Agent", exact: true });
if (premium) {
await expect(toggle).toBeVisible();
if ((await toggle.getAttribute("aria-checked")) !== String(hidden)) {
await toggle.click();
}
await expect(toggle).toHaveAttribute("aria-checked", String(hidden));
} else {
await expect(toggle).toHaveCount(0);
}
await expect(toggle).toHaveAttribute("aria-checked", "true");
await page.keyboard.press("Escape");
await expect(panel).toBeHidden();
await expect(page.getByRole("button", { name: "LiteAdmin", exact: true })).toBeHidden();
const aria = page.getByRole("banner").getByRole("button", { name: "Aria Agent", exact: true });
if (premium && hidden) {
await expect(aria).toBeHidden();
} else {
await expect(aria).toBeVisible();
}
}
/**

View file

@ -1,8 +1,14 @@
import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation";
import { navigateToPage, dismissFeedbackPopup, setAriaHidden } from "../../helpers/navigation";
import { Page } from "../../fixtures/pages";
import { CHAT_MODEL_A, createVirtualKey, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic";
import {
CHAT_MODEL_A,
createVirtualKey,
deleteVirtualKey,
sendChatCompletion,
waitForSpendLog,
} from "../../helpers/traffic";
/**
* Session-grouped pagination (#38060): a page of N rows must render exactly N session rows, a
@ -45,6 +51,57 @@ async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"):
await page.getByRole("option", { name: size, exact: true }).click();
}
test.describe("Logs pagination with Aria in the header", () => {
test.use({ storageState: ADMIN_STORAGE_PATH });
for (const viewport of [
{ width: 1440, height: 900 },
{ width: 390, height: 844 },
]) {
test(`keeps pagination reachable and Aria available at ${viewport.width}px`, async ({ page, request }) => {
await page.setViewportSize(viewport);
const alias = `e2e-aria-header-${uniqueSuffix()}`;
const key = await createVirtualKey(request, { key_alias: alias });
try {
const ids = await Promise.all(
Array.from({ length: 26 }, (_, index) =>
sendChatCompletion(request, {
model: CHAT_MODEL_A,
prompt: `${alias}-${index}`,
apiKey: key.key,
}),
),
);
await Promise.all(ids.map((id) => waitForSpendLog(request, id)));
await openLogs(page);
await setAriaHidden(page, false);
await applyKeyAliasFilter(page, await openFilterDrawer(page), alias);
await setRowsPerPage(page, "25");
await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26");
await visibleTestId(page, "pagination-next").click();
await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 26-26 of 26");
const aria = page.getByRole("banner").getByRole("button", { name: "Aria Agent", exact: true });
await aria.click();
await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible();
await page.keyboard.press("Escape");
await expect(aria).toBeFocused();
await visibleTestId(page, "pagination-first").click();
await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26");
for (const destination of [Page.Usage, Page.LlmPlayground, Page.ApiKeys]) {
await navigateToPage(page, destination);
await aria.click();
await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible();
await page.keyboard.press("Escape");
await expect(aria).toBeFocused();
}
} finally {
await deleteVirtualKey(request, key.token);
}
});
}
});
test.describe("Logs page session-grouped pagination", () => {
test.use({ storageState: ADMIN_STORAGE_PATH });

View file

@ -792,17 +792,23 @@ def _ui_session_cookie(gateway: Gateway, user_id: str) -> dict[str, str]:
def _gateway_session_sign_in(
gateway: Gateway, user_id: str, resource: str | None = None
gateway: Gateway,
user_id: str,
resource: str | None = None,
*,
registration_endpoint: str = "/register",
authorization_endpoint: str = "/authorize",
token_endpoint: str = "/token",
) -> tuple[dict[str, object], str]:
registered: Final = gateway.client.post(
"/register", json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"}
registration_endpoint, json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"}
)
assert registered.status_code in (200, 201), registered.text
client_id: Final = str(_response_object(registered)["client_id"])
pkce: Final = _Pkce(secrets.token_urlsafe(48))
cookies: Final = _ui_session_cookie(gateway, user_id)
started: Final = gateway.client.get(
"/authorize",
authorization_endpoint,
params={
"client_id": client_id,
"redirect_uri": CLIENT_REDIRECT,
@ -823,13 +829,14 @@ def _gateway_session_sign_in(
callback: Final = parse_qs(urlsplit(completed.headers["location"]).query)
assert "code" in callback, completed.headers["location"]
issued: Final = gateway.client.post(
"/token",
token_endpoint,
data={
"grant_type": "authorization_code",
"code": callback["code"][0],
"redirect_uri": CLIENT_REDIRECT,
"client_id": client_id,
"code_verifier": pkce.verifier,
**({} if resource is None else {"resource": resource}),
},
)
assert issued.status_code == 200, issued.text
@ -1292,10 +1299,18 @@ def test_lite_logout_revokes_the_cli_refresh_token_sent_the_way_the_cli_sends_it
assert completion.status_code == 200, completion.text
def _toolset_rpc(gateway: Gateway, bearer: str, name: str, method: str, params: dict[str, object]) -> Outcome:
def _toolset_rpc(
gateway: Gateway,
bearer: str,
name: str,
method: str,
params: dict[str, object],
*,
route: str = "/toolset/{name}/mcp",
) -> Outcome:
def post(rpc_method: str, rpc_params: dict[str, object]) -> httpx.Response:
return gateway.client.post(
f"/toolset/{name}/mcp",
route.format(name=name),
json={"jsonrpc": "2.0", "id": 1, "method": rpc_method, "params": rpc_params},
headers={"Authorization": f"Bearer {bearer}", "Accept": "application/json, text/event-stream"},
)
@ -1306,7 +1321,12 @@ def _toolset_rpc(gateway: Gateway, bearer: str, name: str, method: str, params:
return _outcome_from_rpc(post(method, params))
def test_gateway_session_bearer_of_a_team_member_is_served_the_team_toolset_on_its_route(gateway: Gateway) -> None:
@pytest.mark.parametrize("route", ["/toolset/{name}/mcp", "/{name}/mcp"])
def test_gateway_session_bearer_of_a_team_member_is_served_the_team_toolset_on_its_route(
gateway: Gateway, route: str
) -> None:
from mcp.shared.auth_utils import check_resource_allowed
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "lit6029sess" + uuid.uuid4().hex[:6]
server_id: Final = register_mcp(scenario, peer, alias)
@ -1315,19 +1335,69 @@ def test_gateway_session_bearer_of_a_team_member_is_served_the_team_toolset_on_i
granted_id: Final = create_toolset(scenario, ((server_id, "add"),), toolset_name=granted_name)
create_toolset(scenario, ((server_id, "multiply"),), toolset_name=withheld_name)
member: Final = scenario.member(scenario.team(object_permission={"mcp_toolsets": [granted_id]}))
bearer: Final = _gateway_session_bearer(gateway, member)
path: Final = route.format(name=granted_name)
anonymous: Final = gateway.client.post(
path,
headers=ACCEPT,
json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": dict(INITIALIZE)},
)
assert anonymous.status_code == 401, anonymous.text
metadata_url: Final = anonymous.headers["www-authenticate"].split('resource_metadata="', 1)[1].split('"', 1)[0]
protected: Final = gateway.client.get(metadata_url)
assert protected.status_code == 200, protected.text
protected_body: Final = _response_object(protected)
resource: Final = str(protected_body["resource"])
assert check_resource_allowed(
requested_resource=f"{_base(gateway)}{path}", configured_resource=resource
), protected.text
issuers: Final = TypeAdapter(tuple[str, ...]).validate_python(protected_body["authorization_servers"])
assert len(issuers) == 1, protected.text
issuer: Final = urlsplit(issuers[0])
authorization: Final = gateway.client.get(
f"{issuer.scheme}://{issuer.netloc}/.well-known/oauth-authorization-server{issuer.path}"
)
assert authorization.status_code == 200, authorization.text
authorization_body: Final = _response_object(authorization)
assert authorization_body["issuer"] == issuers[0], authorization.text
token_endpoint: Final = str(authorization_body["token_endpoint"])
issued, client_id = _gateway_session_sign_in(
gateway,
member,
resource,
registration_endpoint=str(authorization_body["registration_endpoint"]),
authorization_endpoint=str(authorization_body["authorization_endpoint"]),
token_endpoint=token_endpoint,
)
bearer: Final = _issued_token(issued)
assert bearer.startswith("llm_session_"), bearer[:16]
listed: Final = _toolset_rpc(gateway, bearer, granted_name, "tools/list", {})
listed: Final = _toolset_rpc(gateway, bearer, granted_name, "tools/list", {}, route=route)
assert listed.tools == (f"{alias}-add",), listed.raw
peer.drain()
called: Final = _toolset_rpc(
gateway, bearer, granted_name, "tools/call", {"name": f"{alias}-add", "arguments": {"a": 4, "b": 5}}
gateway, bearer, granted_name, "tools/call",
{"name": f"{alias}-add", "arguments": {"a": 4, "b": 5}}, route=route
)
assert called.ok and called.text == "9", called.raw
assert len(tool_calls(peer.drain())) == 1
denied: Final = _toolset_rpc(gateway, bearer, withheld_name, "tools/list", {})
denied: Final = _toolset_rpc(gateway, bearer, withheld_name, "tools/list", {}, route=route)
assert denied.status == 403, denied.raw
assert tool_calls(peer.drain()) == ()
rotated: Final = gateway.client.post(
token_endpoint,
data={
"grant_type": "refresh_token",
"refresh_token": str(issued["refresh_token"]),
"client_id": client_id,
"resource": resource,
},
)
assert rotated.status_code == 200, rotated.text
rotated_body: Final = _response_object(rotated)
assert rotated_body["refresh_token"] != issued["refresh_token"], rotated.text
refreshed: Final = _toolset_rpc(
gateway, _issued_token(rotated_body), granted_name, "tools/list", {}, route=route
)
assert refreshed.tools == listed.tools, refreshed.raw
def test_resource_scoped_session_bearer_opens_a_team_toolset_inside_its_server_and_none_outside(

View file

@ -0,0 +1,474 @@
from __future__ import annotations
import asyncio
import json
import linecache
import sys
import threading
import uuid
from collections.abc import AsyncIterable, Awaitable, Callable, Iterable
from concurrent.futures import ThreadPoolExecutor
from types import FrameType
from typing import Final, TypeVar
import pytest
from integration._support.client import object_value
from integration._support.openai_wire import chat_reply, openai_error, responses_reply
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue
import litellm
from litellm import Router
_DEPLOYMENTS: Final = 4
_TARGET: Final = _DEPLOYMENTS - 1
_SYNC_SECONDS: Final = 10
_RATE_LIMITED_KEY: Final = "key-rate-limited"
_BURST_THREADS: Final = 8
_BURST_CALLS: Final = 24
_CHURN_UPSERTS: Final = 120
_UNREPORTED_MODEL_ID: Final = frozenset({"aresponses-stream", "aanthropic-messages"})
_ROUTER_FILE: Final = Router.__init__.__code__.co_filename
_INDEX_LOCK_LINES: Final = frozenset(
number
for number, line in enumerate(linecache.getlines(_ROUTER_FILE), start=1)
if "with self._deployment_index_lock" in line
)
T = TypeVar("T")
def _peer(request: Request) -> Reply:
if request.headers["authorization"] == f"Bearer {_RATE_LIMITED_KEY}":
return openai_error(429)
body: Final = object_value(json.loads(request.body))
stream: Final = body.get("stream") is True
if request.target == "/v1/responses":
return responses_reply(f"resp_{uuid.uuid4().hex}", "gpt-4o-mini", "ok", stream=stream)
return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", "gpt-4o-mini", "ok", stream=stream)
def _deployment(
index: int,
api_base: str,
team: bool,
*,
model: int | None = None,
public_name: str | None = None,
timeout: float = 10,
) -> dict[str, JsonValue]:
team_info: Final[dict[str, JsonValue]] = (
{"team_id": "team-1", "team_public_model_name": public_name or f"public-{index}"} if team else {}
)
return {
"model_name": f"model-{index if model is None else model}",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_base": api_base,
"api_key": f"key-{index}",
"timeout": timeout,
},
"model_info": {"id": f"deployment-{index}", **team_info},
}
def _router(wire: Wire, team: bool = False) -> Router:
return Router(model_list=[_deployment(index, f"{wire.url}/v1", team) for index in range(_DEPLOYMENTS)])
def _read_during_held_write(read: Callable[[], T], write: Callable[[], object]) -> T:
"""Run `read` on its own thread while `write` waits where the list has shrunk but the index maps have not caught
up. The write resumes once `read` returns or reaches the router's index lock, so a read that skips the lock meets
the half-updated list and a read that takes the lock waits for the write"""
reindex: Final = Router._update_deployment_indices_after_removal.__code__
held: Final = threading.Event()
release: Final = threading.Event()
settled: Final = threading.Event()
released_in_time: Final[list[bool]] = []
outcome: Final[list[T]] = []
failure: Final[list[BaseException]] = []
def hold_at_reindex(frame: FrameType, event: str, _arg: object) -> None:
if event == "call" and frame.f_code is reindex and not held.is_set():
held.set()
released_in_time.append(release.wait(timeout=_SYNC_SECONDS))
def run_write() -> None:
previous: Final = sys.getprofile()
sys.setprofile(hold_at_reindex)
try:
write()
finally:
sys.setprofile(previous)
def watch_read(frame: FrameType, event: str, _arg: object) -> Callable[..., object] | None:
if frame.f_code.co_filename != _ROUTER_FILE:
return None
if event == "line" and frame.f_lineno in _INDEX_LOCK_LINES:
settled.set()
return watch_read
def run_read() -> None:
previous: Final = sys.gettrace()
sys.settrace(watch_read)
try:
outcome.append(read())
except BaseException as error:
failure.append(error)
finally:
sys.settrace(previous)
settled.set()
writer: Final = threading.Thread(target=run_write)
reader: Final = threading.Thread(target=run_read)
writer.start()
try:
assert held.wait(timeout=_SYNC_SECONDS), "the write never reached re-indexing"
reader.start()
assert settled.wait(timeout=_SYNC_SECONDS), "the read neither returned nor reached the index lock"
finally:
release.set()
if reader.ident is not None:
reader.join()
writer.join()
assert released_in_time == [True], "the write was not held for the whole read"
if failure:
raise failure[0]
return outcome[0]
def _read_during_held_delete(router: Router, deployment_id: str, read: Callable[[], T]) -> T:
return _read_during_held_write(read, lambda: router.delete_deployment(deployment_id))
def _keys(wire: Wire) -> tuple[str, ...]:
return tuple(request.headers["authorization"] for request in wire.drain())
def _team_metadata(team: bool) -> dict[str, JsonValue]:
return {"user_api_key_team_id": "team-1"} if team else {}
def _completion(router: Router, model: str, team: bool) -> str | None:
response: Final = router.completion(
model=model, messages=[{"role": "user", "content": "hi"}], metadata=_team_metadata(team)
)
return response._hidden_params["model_id"]
def _completion_stream(router: Router, model: str, team: bool) -> str | None:
stream: Final = router.completion(
model=model, messages=[{"role": "user", "content": "hi"}], metadata=_team_metadata(team), stream=True
)
assert isinstance(stream, Iterable), stream
chunks: Final = list(stream)
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok", chunks
return stream._hidden_params["model_id"]
def _acompletion(router: Router, model: str, team: bool) -> str | None:
async def call() -> str | None:
response: Final = await router.acompletion(
model=model, messages=[{"role": "user", "content": "hi"}], metadata=_team_metadata(team)
)
return response._hidden_params["model_id"]
return asyncio.run(call())
def _acompletion_stream(router: Router, model: str, team: bool) -> str | None:
async def call() -> str | None:
stream: Final = await router.acompletion(
model=model, messages=[{"role": "user", "content": "hi"}], metadata=_team_metadata(team), stream=True
)
assert isinstance(stream, AsyncIterable), stream
chunks: Final = [chunk async for chunk in stream]
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok", chunks
return stream._hidden_params["model_id"]
return asyncio.run(call())
def _responses(router: Router, model: str, team: bool) -> str | None:
response: Final = router.responses(model=model, input="hi", metadata=_team_metadata(team))
assert response.output_text == "ok", response
return response._hidden_params["model_id"]
def _aresponses(router: Router, model: str, team: bool) -> str | None:
async def call() -> str | None:
response: Final = await router.aresponses(model=model, input="hi", metadata=_team_metadata(team))
assert response.output_text == "ok", response
return response._hidden_params["model_id"]
return asyncio.run(call())
def _aresponses_stream(router: Router, model: str, team: bool) -> str | None:
async def call() -> None:
stream: Final = await router.aresponses(model=model, input="hi", metadata=_team_metadata(team), stream=True)
assert isinstance(stream, AsyncIterable), stream
events: Final = [event.type async for event in stream]
assert events[-1] == "response.completed", events
asyncio.run(call())
return None
def _aanthropic_messages(router: Router, model: str, team: bool) -> str | None:
async def call() -> None:
response: Final = await router.aanthropic_messages(
model=model, messages=[{"role": "user", "content": "hi"}], max_tokens=16, metadata=_team_metadata(team)
)
assert response["content"][0]["text"] == "ok", response
asyncio.run(call())
return None
_CALLS: Final[dict[str, Callable[[Router, str, bool], str | None]]] = {
"completion": _completion,
"completion-stream": _completion_stream,
"acompletion": _acompletion,
"acompletion-stream": _acompletion_stream,
"responses": _responses,
"aresponses": _aresponses,
"aresponses-stream": _aresponses_stream,
"aanthropic-messages": _aanthropic_messages,
}
@pytest.mark.parametrize("call", list(_CALLS))
@pytest.mark.parametrize("team", [False, True], ids=["model-name", "team-public-name"])
def test_call_reaches_its_deployment_while_another_thread_deletes_one(call: str, team: bool) -> None:
with wire_server(_peer) as wire:
router: Final = _router(wire, team)
served_by: Final = _read_during_held_delete(
router,
"deployment-0",
lambda: _CALLS[call](router, f"public-{_TARGET}" if team else f"model-{_TARGET}", team),
)
keys: Final = _keys(wire)
assert served_by == (None if call in _UNREPORTED_MODEL_ID else f"deployment-{_TARGET}")
assert keys == (f"Bearer key-{_TARGET}",)
assert router.get_model_ids() == [f"deployment-{index}" for index in range(1, _DEPLOYMENTS)]
@pytest.mark.parametrize("deleted", [0, _TARGET - 1], ids=["first", "next-to-target"])
def test_completion_reaches_its_deployment_whichever_earlier_deployment_is_deleted(deleted: int) -> None:
with wire_server(_peer) as wire:
router: Final = _router(wire)
served_by: Final = _read_during_held_delete(
router, f"deployment-{deleted}", lambda: _completion(router, f"model-{_TARGET}", team=False)
)
keys: Final = _keys(wire)
assert served_by == f"deployment-{_TARGET}"
assert keys == (f"Bearer key-{_TARGET}",)
def test_team_completion_reaches_the_remaining_deployment_of_a_shared_public_name_while_its_sibling_is_deleted() -> (
None
):
with wire_server(_peer) as wire:
api_base: Final = f"{wire.url}/v1"
router: Final = Router(
model_list=[
_deployment(0, api_base, team=False),
_deployment(1, api_base, team=False),
_deployment(2, api_base, team=True, public_name="public-shared"),
_deployment(3, api_base, team=True, public_name="public-shared"),
]
)
served_by: Final = _read_during_held_delete(
router, "deployment-2", lambda: _completion(router, "public-shared", team=True)
)
keys: Final = _keys(wire)
assert served_by == "deployment-3"
assert keys == ("Bearer key-3",)
def test_completion_for_the_deployment_being_deleted_fails_cleanly_and_sends_nothing() -> None:
with wire_server(_peer) as wire:
router: Final = _router(wire)
with pytest.raises(litellm.BadRequestError):
_read_during_held_delete(
router,
"deployment-0",
lambda: router.completion(model="model-0", messages=[{"role": "user", "content": "hi"}]),
)
keys: Final = _keys(wire)
assert keys == ()
assert router.get_model_ids() == [f"deployment-{index}" for index in range(1, _DEPLOYMENTS)]
def test_provider_error_during_a_concurrent_delete_comes_from_the_requested_deployment() -> None:
with wire_server(_peer) as wire:
api_base: Final = f"{wire.url}/v1"
rate_limited: Final = _deployment(_TARGET, api_base, team=False)
rate_limited["litellm_params"] = {**object_value(rate_limited["litellm_params"]), "api_key": _RATE_LIMITED_KEY}
router: Final = Router(
model_list=[*(_deployment(index, api_base, team=False) for index in range(_TARGET)), rate_limited],
num_retries=0,
)
with pytest.raises(litellm.RateLimitError):
_read_during_held_delete(
router,
"deployment-0",
lambda: router.completion(model=f"model-{_TARGET}", messages=[{"role": "user", "content": "hi"}]),
)
served_by: Final = _completion(router, "model-1", team=False)
keys: Final = _keys(wire)
assert served_by == "deployment-1"
assert keys == (f"Bearer {_RATE_LIMITED_KEY}", "Bearer key-1")
def test_two_threads_deleting_the_same_deployment_remove_it_once() -> None:
with wire_server(_peer) as wire:
router: Final = _router(wire)
second: Final = _read_during_held_delete(
router, "deployment-0", lambda: router.delete_deployment("deployment-0")
)
served_by: Final = _completion(router, f"model-{_TARGET}", team=False)
keys: Final = _keys(wire)
assert second is None
assert served_by == f"deployment-{_TARGET}"
assert keys == (f"Bearer key-{_TARGET}",)
assert router.get_model_ids() == [f"deployment-{index}" for index in range(1, _DEPLOYMENTS)]
def test_completion_reaches_its_deployment_while_another_thread_upserts_an_earlier_one() -> None:
with wire_server(_peer) as wire:
router: Final = _router(wire)
changed: Final = litellm.Deployment(**_deployment(0, f"{wire.url}/v1", team=False, timeout=11))
served_by: Final = _read_during_held_write(
lambda: _completion(router, f"model-{_TARGET}", team=False), lambda: router.upsert_deployment(changed)
)
keys: Final = _keys(wire)
assert served_by == f"deployment-{_TARGET}"
assert keys == (f"Bearer key-{_TARGET}",)
assert sorted(router.get_model_ids()) == [f"deployment-{index}" for index in range(_DEPLOYMENTS)]
def test_deployment_added_during_a_concurrent_delete_serves_its_own_requests() -> None:
with wire_server(_peer) as wire:
router: Final = _router(wire)
added: Final = litellm.Deployment(**_deployment(9, f"{wire.url}/v1", team=False))
_read_during_held_delete(router, "deployment-0", lambda: router.add_deployment(added))
served: Final = tuple(_completion(router, f"model-{index}", team=False) for index in (9, _TARGET))
keys: Final = _keys(wire)
assert served == ("deployment-9", f"deployment-{_TARGET}")
assert keys == ("Bearer key-9", f"Bearer key-{_TARGET}")
def test_model_list_set_during_a_concurrent_delete_serves_its_own_requests() -> None:
with wire_server(_peer) as wire:
router: Final = _router(wire)
replacement: Final = [_deployment(index, f"{wire.url}/v1", team=False) for index in range(10, 14)]
_read_during_held_delete(router, "deployment-0", lambda: router.set_model_list(replacement))
served: Final = tuple(_completion(router, f"model-{index}", team=False) for index in range(10, 14))
keys: Final = _keys(wire)
assert served == tuple(f"deployment-{index}" for index in range(10, 14))
assert keys == tuple(f"Bearer key-{index}" for index in range(10, 14))
assert router.get_model_ids() == [f"deployment-{index}" for index in range(10, 14)]
def _burst_router(api_base: str) -> Router:
"""Two deployments per model name, so an upsert of one leaves its model group served."""
return Router(
model_list=[
_deployment(index, api_base, team=False, model=index % _DEPLOYMENTS) for index in range(2 * _DEPLOYMENTS)
]
)
def _churn(router: Router, api_base: str, stop: threading.Event) -> int:
"""Upsert the deployments in turn with a changed timeout; each upsert moves one to the end of the list."""
upserts = 0
while upserts < _CHURN_UPSERTS and not stop.is_set():
index, round_number = upserts % (2 * _DEPLOYMENTS), upserts // (2 * _DEPLOYMENTS)
router.upsert_deployment(
litellm.Deployment(
**_deployment(index, api_base, team=False, model=index % _DEPLOYMENTS, timeout=11 + round_number % 2)
)
)
upserts += 1
return upserts
def _served_keys(wire: Wire) -> dict[str, str]:
"""Which key reached the upstream for each request marker."""
return {
object_value(json.loads(request.body))["messages"][0]["content"]: request.headers["authorization"]
for request in wire.drain()
}
def _assert_each_request_used_its_models_deployment(served: list[tuple[str, str | None]], keys: dict[str, str]) -> None:
for marker, model_id in served:
model: Final = int(marker.split(":")[0])
assert model_id in (f"deployment-{model}", f"deployment-{model + _DEPLOYMENTS}"), (marker, model_id)
assert keys[marker] == f"Bearer key-{str(model_id).removeprefix('deployment-')}", (marker, keys[marker])
assert len(keys) == len(served)
def test_thread_burst_during_upsert_churn_sends_every_request_with_its_deployments_key() -> None:
with wire_server(_peer) as wire:
api_base: Final = f"{wire.url}/v1"
router: Final = _burst_router(api_base)
stop: Final = threading.Event()
def call(number: int) -> tuple[str, str | None]:
marker: Final = f"{number % _DEPLOYMENTS}:{number}"
response: Final = router.completion(
model=f"model-{number % _DEPLOYMENTS}", messages=[{"role": "user", "content": marker}]
)
return marker, response._hidden_params["model_id"]
with ThreadPoolExecutor(max_workers=_BURST_THREADS + 1) as pool:
churn: Final = pool.submit(_churn, router, api_base, stop)
served: Final = list(pool.map(call, range(_BURST_THREADS * _BURST_CALLS)))
stop.set()
upserts: Final = churn.result()
keys: Final = _served_keys(wire)
assert upserts > 0
_assert_each_request_used_its_models_deployment(served, keys)
assert sorted(router.get_model_ids()) == sorted(f"deployment-{index}" for index in range(2 * _DEPLOYMENTS))
async def _gathered(calls: Iterable[Awaitable[tuple[str, str | None]]]) -> list[tuple[str, str | None]]:
return list(await asyncio.gather(*calls))
def test_async_burst_during_upsert_churn_sends_every_request_with_its_deployments_key() -> None:
with wire_server(_peer) as wire:
api_base: Final = f"{wire.url}/v1"
router: Final = _burst_router(api_base)
stop: Final = threading.Event()
async def call(number: int) -> tuple[str, str | None]:
marker: Final = f"{number % _DEPLOYMENTS}:{number}"
response: Final = await router.acompletion(
model=f"model-{number % _DEPLOYMENTS}", messages=[{"role": "user", "content": marker}]
)
return marker, response._hidden_params["model_id"]
with ThreadPoolExecutor(max_workers=1) as pool:
churn: Final = pool.submit(_churn, router, api_base, stop)
served: Final = asyncio.run(_gathered(call(number) for number in range(_BURST_THREADS * _BURST_CALLS)))
stop.set()
upserts: Final = churn.result()
keys: Final = _served_keys(wire)
assert upserts > 0
_assert_each_request_used_its_models_deployment(served, keys)
assert sorted(router.get_model_ids()) == sorted(f"deployment-{index}" for index in range(2 * _DEPLOYMENTS))

View file

@ -1,5 +1,9 @@
from collections.abc import Iterator
from pathlib import Path
from typing import Final
import respx
import asyncio, httpx, importlib, os
import pytest
@ -8,6 +12,7 @@ from litellm import CustomLLM
from litellm.litellm_core_utils import get_llm_provider_logic
from litellm.litellm_core_utils.get_llm_provider_logic import (
get_llm_provider,
resolve_model_and_provider_for_metadata,
get_provider_affinity_header,
inferred_provider,
is_registered_custom_provider,
@ -715,3 +720,47 @@ def test_opencode_affinity_resolution_preserves_existing_transport(prefix: str)
model, provider, _, destination = get_llm_provider(model=f"{prefix}/test-model", api_base=api_base)
assert (model, provider, destination) == ("test-model", prefix, api_base)
assert get_provider_affinity_header(destination) == "x-opencode-session"
@pytest.fixture
def chatgpt_login_tripwire(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Iterator[None]:
blocked: Final = tmp_path / "blocked"
blocked.write_text("ordinary file")
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(blocked / "unused"))
monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False)
with respx.mock as transport:
yield
assert len(transport.calls) == 0
@pytest.mark.parametrize(
("model", "custom_llm_provider", "expected"),
[
("chatgpt/gpt-5.4", None, ("gpt-5.4", "chatgpt")),
("gpt-5.4", "chatgpt", ("gpt-5.4", "chatgpt")),
("github_copilot/gpt-4o", None, ("gpt-4o", "github_copilot")),
("openai/gpt-5.4", None, ("gpt-5.4", "openai")),
],
)
def test_resolve_model_and_provider_for_metadata_never_logs_in(
chatgpt_login_tripwire: None, model: str, custom_llm_provider: str | None, expected: tuple[str, str]
) -> None:
assert resolve_model_and_provider_for_metadata(model, custom_llm_provider) == expected
def test_metadata_callers_never_log_in_for_a_declared_chatgpt_deployment(
chatgpt_login_tripwire: None, monkeypatch: pytest.MonkeyPatch
) -> None:
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
from litellm.proxy.spend_tracking.savings import _resolve_model
from litellm.router import Router
from litellm.types.router import LiteLLM_Params
from litellm.utils import supports_native_streaming
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
assert supports_native_streaming("chatgpt/gpt-5.4", None) is True
assert get_api_base("chatgpt/gpt-5.4", {"model": "chatgpt/gpt-5.4"}) == "https://chatgpt.com/backend-api/codex"
identity = _resolve_model("chatgpt/gpt-5.4", None)
assert identity is not None and identity.provider == "chatgpt"
assert Router._model_group_llm_provider(LiteLLM_Params(model="chatgpt/gpt-5.4")) == ("gpt-5.4", "chatgpt")

View file

@ -0,0 +1,638 @@
from unittest.mock import MagicMock
import pytest
from litellm.exceptions import AuthenticationError
from litellm.llms.chatgpt.chat.transformation import ChatGPTConfig
from litellm.llms.chatgpt.common_utils import GetAccessTokenError
GATEWAY_BASE = "https://chatgpt.example/backend-api/codex"
def _gateway_authenticator() -> MagicMock:
authenticator = MagicMock()
authenticator.get_api_base.return_value = GATEWAY_BASE
authenticator.get_access_token.return_value = "gateway-access-token"
return authenticator
def _provider_info(
config: ChatGPTConfig, api_base: str | None, api_key: str | None
) -> tuple[str | None, str | None, str]:
return config._get_openai_compatible_provider_info(
model="gpt-5.4", api_base=api_base, api_key=api_key, custom_llm_provider="chatgpt"
)
class TestChatGPTChatProviderInfo:
def test_configured_key_uses_configured_base_without_gateway_login(self):
authenticator = _gateway_authenticator()
config = ChatGPTConfig(authenticator=authenticator)
assert _provider_info(config, "http://127.0.0.1:9999", "configured-key") == (
"http://127.0.0.1:9999",
"configured-key",
"chatgpt",
)
authenticator.get_access_token.assert_not_called()
def test_configured_key_without_base_uses_gateway_base(self):
config = ChatGPTConfig(authenticator=_gateway_authenticator())
assert _provider_info(config, None, "configured-key") == (GATEWAY_BASE, "configured-key", "chatgpt")
def test_gateway_login_is_never_sent_to_a_custom_base(self):
config = ChatGPTConfig(authenticator=_gateway_authenticator())
assert _provider_info(config, "http://attacker.example", None) == (
GATEWAY_BASE,
"gateway-access-token",
"chatgpt",
)
def test_missing_gateway_login_raises_authentication_error(self):
authenticator = _gateway_authenticator()
authenticator.get_access_token.side_effect = GetAccessTokenError(status_code=401, message="not logged in")
config = ChatGPTConfig(authenticator=authenticator)
with pytest.raises(AuthenticationError, match="not logged in"):
_provider_info(config, None, None)
import asyncio
import json
from pathlib import Path
from types import MappingProxyType
from typing import Final
import httpx
import respx
import litellm
from litellm.integrations.dotprompt.dotprompt_manager import DotpromptManager
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.llms.chatgpt.chat.transformation import ChatGPTConfig
from litellm.types.router import GenericLiteLLMParams
def _write_auth_record(path: Path, token: str, account_id: str) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(
json.dumps(
{
"access_token": token,
"account_id": account_id,
"expires_at": 4102444800,
}
)
)
@pytest.fixture
def auth_files(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path, Path]:
monkeypatch.setattr("time.time", lambda: 1700000000)
default_dir: Final = tmp_path / "default"
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(default_dir))
monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False)
_write_auth_record(default_dir / "auth.json", "token-default", "acct-default")
file_a: Final = tmp_path / "account-a" / "auth.json"
file_b: Final = tmp_path / "account-b" / "auth.json"
_write_auth_record(file_a, "token-a", "acct-a")
_write_auth_record(file_b, "token-b", "acct-b")
return file_a, file_b
class TestChatGPTMultiAccountChat:
def test_validate_environment_uses_per_model_auth_file(self, auth_files: tuple[Path, Path]) -> None:
file_a, file_b = auth_files
config: Final = ChatGPTConfig()
headers_a: Final = config.validate_environment(
headers={},
model="gpt-6.1-sol",
messages=[],
optional_params={},
litellm_params={"chatgpt_auth_file": str(file_a)},
api_key="token-default",
)
headers_b: Final = config.validate_environment(
headers={},
model="gpt-6.1-sol",
messages=[],
optional_params={},
litellm_params={"chatgpt_auth_file": str(file_b)},
api_key="token-default",
)
assert headers_a["Authorization"] == "Bearer token-a"
assert headers_a["ChatGPT-Account-Id"] == "acct-a"
assert headers_b["Authorization"] == "Bearer token-b"
assert headers_b["ChatGPT-Account-Id"] == "acct-b"
def test_validate_environment_defaults_to_global_auth_file(self, auth_files: tuple[Path, Path]) -> None:
config: Final = ChatGPTConfig()
headers: Final = config.validate_environment(
headers={},
model="gpt-6.1-sol",
messages=[],
optional_params={},
litellm_params={},
api_key="token-default",
)
assert headers["Authorization"] == "Bearer token-default"
assert headers["ChatGPT-Account-Id"] == "acct-default"
def test_provider_info_uses_per_model_auth_file(self, auth_files: tuple[Path, Path]) -> None:
file_a, _ = auth_files
config: Final = ChatGPTConfig()
_, per_model_key, _ = config._get_openai_compatible_provider_info(
"gpt-6.1-sol", None, None, "chatgpt", {"chatgpt_auth_file": str(file_a)}
)
_, default_key, _ = config._get_openai_compatible_provider_info("gpt-6.1-sol", None, None, "chatgpt")
assert per_model_key == "token-a"
assert default_key == "token-default"
def test_get_llm_provider_threads_auth_file(self, auth_files: tuple[Path, Path]) -> None:
file_a, _ = auth_files
model, provider, dynamic_api_key, _ = get_llm_provider(
model="chatgpt/gpt-6.1-sol",
litellm_params=GenericLiteLLMParams(chatgpt_auth_file=str(file_a)),
)
assert model == "gpt-6.1-sol"
assert provider == "chatgpt"
assert dynamic_api_key == "token-a"
def _subscription_response(request: httpx.Request) -> httpx.Response:
account: Final = request.headers["ChatGPT-Account-Id"]
assert request.headers["Authorization"] in (
f"Bearer token-{account.removeprefix('acct-')}",
f"Bearer {_client_token(account)}",
)
assert "chatgpt_auth_file" not in request.content.decode()
message: Final = {
"id": "msg_test",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": account, "annotations": []}],
}
response: Final = {
"id": "resp_test",
"object": "response",
"created_at": 1700000000,
"status": "completed",
"model": "gpt-6.1-sol",
"output": [message],
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
}
events: Final = (
{"type": "response.created", "response": {**response, "status": "in_progress", "output": []}},
{
"type": "response.output_text.delta",
"item_id": "msg_test",
"output_index": 0,
"content_index": 0,
"delta": account,
},
{"type": "response.output_item.done", "output_index": 0, "item": message},
{"type": "response.completed", "response": response},
)
return httpx.Response(
200,
text="".join(f"data: {json.dumps(event)}\n\n" for event in events),
headers={"content-type": "text/event-stream"},
)
@pytest.mark.parametrize("blocked_default", [False, True])
@pytest.mark.parametrize("with_prompt", [False, True])
@pytest.mark.parametrize("endpoint", ["chat", "responses"])
@pytest.mark.parametrize("through_router", [False, True])
@pytest.mark.asyncio
async def test_requests_use_selected_subscription(
auth_files: tuple[Path, Path],
endpoint: str,
through_router: bool,
with_prompt: bool,
blocked_default: bool,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
default_dir: Final = auth_files[0].parent / "missing-default"
if blocked_default:
default_dir.write_text("ordinary file")
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(default_dir / "unused"))
if with_prompt:
monkeypatch.setattr(
litellm,
"callbacks",
[
DotpromptManager(
prompt_id="account-check",
prompt_data={"content": "Saved account prompt", "metadata": {"model": "chatgpt/gpt-6.1-sol"}},
)
],
)
router: Final = (
litellm.Router(
model_list=[
{
"model_name": f"account-{index}",
"litellm_params": {"model": "chatgpt/gpt-6.1-sol", "chatgpt_auth_file": str(path)},
}
for index, path in enumerate(auth_files)
]
)
if through_router
else None
)
async def request(index: int, path: Path) -> str:
params: Final = MappingProxyType(
{
"model": f"account-{index}" if router else "chatgpt/gpt-6.1-sol",
**({"chatgpt_auth_file": str(path)} if router is None else {}),
**({"prompt_id": "account-check"} if with_prompt else {}),
}
)
if endpoint == "chat":
result: Final = await (router.acompletion if router else litellm.acompletion)(
**params, messages=[{"role": "user", "content": "Hello"}]
)
return result.choices[0].message.content
response: Final = await (router.aresponses if router else litellm.aresponses)(
**params, input=[{"role": "user", "content": "Hello"}]
)
return response.output[0].content[0].text
with respx.mock as transport:
transport.post("https://chatgpt.com/backend-api/codex/responses").mock(side_effect=_subscription_response)
results: Final = await asyncio.gather(*(request(index, path) for index, path in enumerate(auth_files)))
assert results == ["acct-a", "acct-b"]
if with_prompt:
assert all("Saved account prompt" in call.request.content.decode() for call in transport.calls)
def _client_token(account: str) -> str:
import base64
payload: Final = (
base64.urlsafe_b64encode(json.dumps({"https://api.openai.com/auth": {"chatgpt_account_id": account}}).encode())
.rstrip(b"=")
.decode()
)
return f"header.{payload}.signature"
@pytest.mark.parametrize("endpoint", ["chat", "responses"])
@pytest.mark.parametrize("configured_login", [False, True])
@pytest.mark.parametrize("with_prompt", [False, True])
@pytest.mark.parametrize("stream", [False, True])
@pytest.mark.asyncio
async def test_client_login_keeps_its_account_across_gateway_deployments(
auth_files: tuple[Path, Path],
endpoint: str,
configured_login: bool,
with_prompt: bool,
stream: bool,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(auth_files[0].parent / "no-gateway-login"))
if with_prompt:
monkeypatch.setattr(
litellm,
"callbacks",
[
DotpromptManager(
prompt_id="client-account-check",
prompt_data={"content": "Saved client prompt", "metadata": {"model": "chatgpt/gpt-6.1-sol"}},
)
],
)
router: Final = litellm.Router(
model_list=[
{
"model_name": "subscriptions",
"litellm_params": {
"model": "chatgpt/gpt-6.1-sol",
**({"chatgpt_auth_file": str(path)} if configured_login else {}),
},
}
for path in auth_files
],
num_retries=0,
)
def reply(request: httpx.Request) -> httpx.Response:
account: Final = request.headers["ChatGPT-Account-Id"]
assert request.headers["Authorization"] == f"Bearer {_client_token(account)}"
return _subscription_response(
httpx.Request(
"POST",
request.url,
content=request.content,
headers={
**{key: value for key, value in request.headers.items() if key.lower() != "authorization"},
"Authorization": f"Bearer token-{account.removeprefix('acct-')}",
},
)
)
async def call(account: str) -> str:
params: Final = MappingProxyType(
{
"model": "subscriptions",
"stream": stream,
**({"prompt_id": "client-account-check"} if with_prompt else {}),
"provider_specific_header": {
"custom_llm_provider": "chatgpt",
"extra_headers": {"Authorization": f"Bearer {_client_token(account)}"},
},
}
)
response: Final = await (
router.acompletion(**params, messages=[{"role": "user", "content": "Hello"}])
if endpoint == "chat"
else router.aresponses(**params, input=[{"role": "user", "content": "Hello"}])
)
if stream:
chunks: Final = [part async for part in response]
return (
"".join(part.choices[0].delta.content or "" for part in chunks)
if endpoint == "chat"
else "".join(
getattr(part, "delta", "")
for part in chunks
if getattr(part, "type", "") == "response.output_text.delta"
)
)
return response.choices[0].message.content if endpoint == "chat" else response.output[0].content[0].text
with respx.mock as transport:
transport.post("https://chatgpt.com/backend-api/codex/responses").mock(side_effect=reply)
results: Final = await asyncio.gather(call("acct-client-a"), call("acct-client-b"))
assert results == ["acct-client-a", "acct-client-b"]
assert len(transport.calls) == 2
@pytest.mark.asyncio
async def test_one_client_limit_does_not_cool_down_other_subscriptions(
auth_files: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
router: Final = litellm.Router(
model_list=[
{"model_name": "pool", "litellm_params": {"model": "chatgpt/gpt-6.1-sol", "chatgpt_auth_file": str(path)}}
for path in auth_files
],
num_retries=0,
allowed_fails=0,
cooldown_time=60,
)
headers: Final = {
"custom_llm_provider": "chatgpt",
"extra_headers": {"Authorization": f"Bearer {_client_token('acct-limited')}"},
}
with respx.mock as transport:
route: Final = transport.post("https://chatgpt.com/backend-api/codex/responses")
route.respond(
429,
headers={"retry-after": "60"},
json={"error": {"message": "Usage limit reached", "type": "rate_limit_error"}},
)
for _ in range(2):
with pytest.raises(litellm.RateLimitError):
await router.aresponses(model="pool", input="Hello", provider_specific_header=headers)
await asyncio.sleep(0)
route.mock(side_effect=_subscription_response)
response: Final = await router.aresponses(model="pool", input="Hello")
assert response.output[0].content[0].text in ("acct-a", "acct-b")
@pytest.mark.parametrize("endpoint", ["chat", "responses"])
@pytest.mark.parametrize("async_call", [False, True])
@pytest.mark.parametrize("direct_headers", [False, True])
@pytest.mark.parametrize("configured_key", [False, True])
@pytest.mark.parametrize(
"target", ["https://example.com/v1", "http://chatgpt.com/backend-api/codex", "https://chatgpt.com.example.com/v1"]
)
@pytest.mark.asyncio
async def test_client_login_cannot_be_redirected_to_another_host(
endpoint: str, target: str, async_call: bool, direct_headers: bool, configured_key: bool
) -> None:
headers: Final = {"Authorization": f"Bearer {_client_token('acct-client')}"}
credential: Final = (
{"extra_headers": headers}
if direct_headers
else {"provider_specific_header": {"custom_llm_provider": "chatgpt", "extra_headers": headers}}
)
params: Final = MappingProxyType(
{
"model": "chatgpt/gpt-6.1-sol",
"api_base": target,
**({"api_key": "configured-test-key"} if configured_key else {}),
**credential,
}
)
async def request() -> None:
if endpoint == "chat":
if async_call:
await litellm.acompletion(**params, messages=[{"role": "user", "content": "Hello"}])
else:
litellm.completion(**params, messages=[{"role": "user", "content": "Hello"}])
elif async_call:
await litellm.aresponses(**params, input="Hello")
else:
litellm.responses(**params, input="Hello")
with respx.mock as transport:
with pytest.raises(litellm.BadRequestError, match="official HTTPS ChatGPT endpoint"):
await request()
assert len(transport.calls) == 0
@pytest.mark.asyncio
async def test_cross_provider_fallback_does_not_receive_the_chatgpt_login(
auth_files: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
router: Final = litellm.Router(
model_list=[
{
"model_name": "subscription",
"litellm_params": {"model": "chatgpt/gpt-6.1-sol", "chatgpt_auth_file": str(auth_files[0])},
},
{
"model_name": "api-fallback",
"litellm_params": {"model": "openai/gpt-6.1-sol", "api_key": "sk-example-fallback"},
},
],
fallbacks=[{"subscription": ["api-fallback"]}],
num_retries=0,
)
with respx.mock as transport:
transport.post("https://chatgpt.com/backend-api/codex/responses").respond(
429, json={"error": {"message": "Usage limit reached", "type": "rate_limit_error"}}
)
fallback: Final = transport.post("https://api.openai.com/v1/responses").respond(
200,
json={
"id": "resp_fallback",
"object": "response",
"created_at": 1700000000,
"status": "completed",
"model": "gpt-6.1-sol",
"output": [
{
"id": "msg_fallback",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "fallback", "annotations": []}],
}
],
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
},
)
response: Final = await router.aresponses(
model="subscription",
input="Hello",
provider_specific_header={
"custom_llm_provider": "chatgpt",
"extra_headers": {
"Authorization": f"Bearer {_client_token('acct-client')}",
"ChatGPT-Account-Id": "acct-client",
},
},
)
assert response.output[0].content[0].text == "fallback"
assert fallback.calls.last.request.headers["Authorization"] == "Bearer sk-example-fallback"
assert "ChatGPT-Account-Id" not in fallback.calls.last.request.headers
@pytest.mark.parametrize("endpoint", ["chat", "responses"])
@pytest.mark.parametrize("async_call", [False, True])
@pytest.mark.parametrize("with_prompt", [False, True])
@pytest.mark.asyncio
async def test_sdk_extra_headers_use_client_login_without_gateway_auth(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, endpoint: str, async_call: bool, with_prompt: bool
) -> None:
blocked: Final = tmp_path / "blocked"
blocked.write_text("ordinary file")
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(blocked / "unused"))
monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
headers: Final = {"Authorization": f"Bearer {_client_token('acct-client-a')}"}
prompt_params: Final = {"prompt_id": "sdk-client"} if with_prompt else {}
if with_prompt:
monkeypatch.setattr(
litellm,
"callbacks",
[
DotpromptManager(
prompt_id="sdk-client",
prompt_data={"content": "Hello", "metadata": {"model": "chatgpt/gpt-6-luna"}},
)
],
)
with respx.mock as transport:
transport.post("https://chatgpt.com/backend-api/codex/responses").mock(side_effect=_subscription_response)
if endpoint == "chat":
params: Final = {
"model": "chatgpt/gpt-6.1-sol",
"messages": [{"role": "user", "content": "Hello"}],
"extra_headers": headers,
**prompt_params,
}
result: Final = await litellm.acompletion(**params) if async_call else litellm.completion(**params)
assert result.choices[0].message.content == "acct-client-a"
else:
response: Final = (
await litellm.aresponses(
model="chatgpt/gpt-6.1-sol", input="Hello", extra_headers=headers, **prompt_params
)
if async_call
else litellm.responses(
model="chatgpt/gpt-6.1-sol", input="Hello", extra_headers=headers, **prompt_params
)
)
assert response.output[0].content[0].text == "acct-client-a"
assert transport.calls.last.request.headers["Authorization"] == headers["Authorization"]
assert len(transport.calls) == 1
@pytest.mark.parametrize("endpoint", ["chat", "responses"])
@pytest.mark.parametrize("sync", [False, True])
@pytest.mark.parametrize("credential_source", ["default", "file", "api_key"])
@pytest.mark.parametrize("authorization_header", [None, "Authorization", "authorization"])
@pytest.mark.asyncio
async def test_requests_preserve_gateway_token_and_selected_workspace(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
endpoint: str,
sync: bool,
credential_source: str,
authorization_header: str | None,
) -> None:
monkeypatch.setattr("time.time", lambda: 1700000000)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("CHATGPT_AUTH_FILE", str(tmp_path / "default.json"))
token: Final = _client_token("acct-token-default")
auth_file: Final = tmp_path / ("selected.json" if credential_source == "file" else "default.json")
_write_auth_record(auth_file, token, "acct-saved-workspace")
configured_token: Final = _client_token("acct-configured")
expected_token: Final = configured_token if credential_source == "api_key" else token
expected_account: Final = "acct-configured" if credential_source == "api_key" else "acct-saved-workspace"
def reply(request: httpx.Request) -> httpx.Response:
assert request.headers.get_list("authorization") == [f"Bearer {expected_token}"]
assert request.headers.get_list("chatgpt-account-id") == [expected_account]
assert request.headers["originator"] == "codex_cli_rs"
return _subscription_response(
httpx.Request(
"POST",
request.url,
content=request.content,
headers={
"Authorization": f"Bearer token-{expected_account.removeprefix('acct-')}",
"ChatGPT-Account-Id": expected_account,
},
)
)
params: Final = MappingProxyType(
{
"model": "chatgpt/gpt-6.1-sol",
**({"chatgpt_auth_file": str(auth_file)} if credential_source == "file" else {}),
**({"api_key": configured_token} if credential_source == "api_key" else {}),
"extra_headers": {
"originator": "codex_cli_rs",
**({authorization_header: "Bearer sk-unrelated-example"} if authorization_header else {}),
},
}
)
with respx.mock as transport:
transport.post("https://chatgpt.com/backend-api/codex/responses").mock(side_effect=reply)
if endpoint == "chat":
result: Final = (
await asyncio.to_thread(litellm.completion, **params, messages=[{"role": "user", "content": "Hello"}])
if sync
else await litellm.acompletion(**params, messages=[{"role": "user", "content": "Hello"}])
)
assert result.choices[0].message.content == expected_account
else:
response: Final = (
await asyncio.to_thread(litellm.responses, **params, input=[{"role": "user", "content": "Hello"}])
if sync
else await litellm.aresponses(**params, input=[{"role": "user", "content": "Hello"}])
)
assert response.output[0].content[0].text == expected_account

View file

@ -4,8 +4,10 @@ Tests for ChatGPT subscription Responses API transformation
Source: litellm/llms/chatgpt/responses/transformation.py
"""
import base64
import json
from collections.abc import Generator
from pathlib import Path
from typing import Final
from unittest.mock import MagicMock, patch
@ -13,6 +15,7 @@ import httpx
import pytest
import litellm
from litellm.llms.chatgpt.authenticator import Authenticator
from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig
from litellm.llms.openai.common_utils import OpenAIError
from litellm.main import responses_api_bridge_check
@ -395,3 +398,165 @@ class TestChatGPTResponsesAPITransformation:
assert exc_info.value.message == "ChatGPT upstream failed"
assert exc_info.value.status_code == 502
def _fake_chatgpt_oauth_token(account_id: str) -> str:
header = base64.urlsafe_b64encode(b'{"alg":"RS256"}').rstrip(b"=").decode()
payload = (
base64.urlsafe_b64encode(
json.dumps({"https://api.openai.com/auth": {"chatgpt_account_id": account_id}}).encode()
)
.rstrip(b"=")
.decode()
)
return f"{header}.{payload}.fake-signature"
def _gateway_authenticator() -> MagicMock:
authenticator = MagicMock()
authenticator.get_access_token.return_value = "gateway-access-token"
authenticator.get_account_id.return_value = "acct-gateway"
return authenticator
def _config_with(authenticator: MagicMock) -> ChatGPTResponsesAPIConfig:
return ChatGPTResponsesAPIConfig(authenticator=authenticator)
class TestChatGPTClientCredentialForwarding:
def test_client_chatgpt_oauth_wins_over_configured_key_and_gateway_login(self):
authenticator = _gateway_authenticator()
token = _fake_chatgpt_oauth_token("acct-alice")
headers = _config_with(authenticator).validate_environment(
headers={"Authorization": f"Bearer {token}", "originator": "codex_cli_rs"},
model="gpt-5.4",
litellm_params=GenericLiteLLMParams(api_key="configured-gateway-key", litellm_session_id="session-1"),
)
assert headers["Authorization"] == f"Bearer {token}"
assert headers["ChatGPT-Account-Id"] == "acct-alice"
assert headers["originator"] == "codex_cli_rs"
assert headers["accept"] == "text/event-stream"
assert [name for name in headers if name.lower() == "authorization"] == ["Authorization"]
authenticator.get_access_token.assert_not_called()
def test_explicit_account_header_wins_over_jwt_claim(self):
token = _fake_chatgpt_oauth_token("acct-from-jwt")
headers = _config_with(_gateway_authenticator()).validate_environment(
headers={"authorization": f"Bearer {token}", "chatgpt-account-id": "acct-from-header"},
model="gpt-5.4",
litellm_params=GenericLiteLLMParams(),
)
assert headers["ChatGPT-Account-Id"] == "acct-from-header"
assert [name for name in headers if name.lower() == "chatgpt-account-id"] == ["ChatGPT-Account-Id"]
@pytest.mark.parametrize(
"bearer", ["Bearer sk-proj-fake-platform-key", "Bearer not-a-jwt", "Bearer a.bm90LWpzb24.c"]
)
def test_non_chatgpt_bearer_falls_back_to_configured_key(self, bearer, tmp_path):
authenticator = Authenticator(auth_file=str(tmp_path / "missing-default.json"))
headers = ChatGPTResponsesAPIConfig(authenticator=authenticator).validate_environment(
headers={"Authorization": bearer},
model="gpt-5.4",
litellm_params=GenericLiteLLMParams(api_key="configured-gateway-key"),
)
assert "ChatGPT-Account-Id" not in headers
assert headers["Authorization"] == "Bearer configured-gateway-key"
def test_configured_key_is_used_without_client_credential(self):
authenticator = _gateway_authenticator()
headers = _config_with(authenticator).validate_environment(
headers={},
model="gpt-5.4",
litellm_params=GenericLiteLLMParams(api_key="configured-gateway-key"),
)
assert headers["Authorization"] == "Bearer configured-gateway-key"
authenticator.get_access_token.assert_not_called()
def test_gateway_login_is_used_without_client_credential_or_configured_key(self):
headers = _config_with(_gateway_authenticator()).validate_environment(
headers={},
model="gpt-5.4",
litellm_params=GenericLiteLLMParams(),
)
assert headers["Authorization"] == "Bearer gateway-access-token"
assert headers["ChatGPT-Account-Id"] == "acct-gateway"
def test_two_clients_never_share_a_credential(self):
config = _config_with(_gateway_authenticator())
alice = _fake_chatgpt_oauth_token("acct-alice")
bob = _fake_chatgpt_oauth_token("acct-bob")
forwarded = [
config.validate_environment(
headers={"Authorization": f"Bearer {token}"}, model="gpt-5.4", litellm_params=GenericLiteLLMParams()
)
for token in (alice, bob, alice)
]
assert [h["Authorization"] for h in forwarded] == [f"Bearer {alice}", f"Bearer {bob}", f"Bearer {alice}"]
assert [h["ChatGPT-Account-Id"] for h in forwarded] == ["acct-alice", "acct-bob", "acct-alice"]
class TestChatGPTResponsesMultiAccount:
@pytest.fixture(autouse=True)
def fixed_clock(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr("time.time", lambda: 1700000000)
@staticmethod
def _write_auth_record(path: Path, token: str, account_id: str) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(
json.dumps(
{
"access_token": token,
"account_id": account_id,
"expires_at": 4102444800,
}
)
)
def test_validate_environment_uses_per_model_auth_file(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
default_dir: Final = tmp_path / "default"
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(default_dir))
monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False)
self._write_auth_record(default_dir / "auth.json", "token-default", "acct-default")
file_a: Final = tmp_path / "account-a" / "auth.json"
file_b: Final = tmp_path / "account-b" / "auth.json"
self._write_auth_record(file_a, "token-a", "acct-a")
self._write_auth_record(file_b, "token-b", "acct-b")
config: Final = ChatGPTResponsesAPIConfig()
headers_a: Final = config.validate_environment(
headers={},
model="gpt-6.1-sol",
litellm_params=GenericLiteLLMParams(chatgpt_auth_file=str(file_a)),
)
headers_b: Final = config.validate_environment(
headers={},
model="gpt-6.1-sol",
litellm_params=GenericLiteLLMParams(chatgpt_auth_file=str(file_b)),
)
headers_default: Final = config.validate_environment(
headers={},
model="gpt-6.1-sol",
litellm_params=GenericLiteLLMParams(),
)
assert headers_a["Authorization"] == "Bearer token-a"
assert headers_a["ChatGPT-Account-Id"] == "acct-a"
assert headers_b["Authorization"] == "Bearer token-b"
assert headers_b["ChatGPT-Account-Id"] == "acct-b"
assert headers_default["Authorization"] == "Bearer token-default"
assert headers_default["ChatGPT-Account-Id"] == "acct-default"

View file

@ -2,17 +2,24 @@ import base64
import json
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Final
from unittest.mock import MagicMock, mock_open, patch
import pytest
import httpx
import respx
import litellm
from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS, HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS
from litellm.llms.chatgpt.authenticator import (
TOKEN_REFRESH_TIMEOUT_SECONDS,
Authenticator,
get_cached_authenticator,
get_chatgpt_auth_file,
)
from litellm.llms.chatgpt.common_utils import GetAccessTokenError
from litellm.types.router import GenericLiteLLMParams
def _make_jwt(payload: dict) -> str:
@ -161,3 +168,144 @@ class TestChatGPTAuthenticator:
assert account_id == "acct-123"
mock_write.assert_called_once()
assert mock_write.call_args[0][0]["account_id"] == "acct-123"
def test_get_account_id_prefers_stored_value_without_rewriting_file(self, authenticator):
id_token = _make_jwt({"https://api.openai.com/auth": {"chatgpt_account_id": "acct-from-token"}})
auth_data = json.dumps({"account_id": "acct-stored", "id_token": id_token})
with (
patch("builtins.open", mock_open(read_data=auth_data)),
patch.object(authenticator, "_write_auth_file") as mock_write,
):
assert authenticator.get_account_id() == "acct-stored"
mock_write.assert_not_called()
def test_get_account_id_without_login_returns_none(self, authenticator):
with (
patch("builtins.open", mock_open(read_data="{}")),
patch.object(authenticator, "_write_auth_file") as mock_write,
):
assert authenticator.get_account_id() is None
mock_write.assert_not_called()
def _write_auth_record(path: Path, token: str, account_id: str) -> None:
path.write_text(
json.dumps(
{
"access_token": token,
"account_id": account_id,
"expires_at": 4102444800,
}
)
)
class TestChatGPTMultiAccountAuthenticator:
@pytest.fixture(autouse=True)
def fixed_clock(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr("time.time", lambda: 1700000000)
def test_explicit_auth_file_overrides_env(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path / "env-dir"))
monkeypatch.setenv("CHATGPT_AUTH_FILE", "env.json")
custom_file: Final = tmp_path / "account-a" / "auth.json"
authenticator: Final = Authenticator(auth_file=str(custom_file))
assert authenticator.auth_file == str(custom_file)
assert authenticator.token_dir == str(custom_file.parent)
assert not custom_file.parent.exists()
authenticator._write_auth_file({"access_token": "token-a"})
assert json.loads(custom_file.read_text())["access_token"] == "token-a"
def test_default_auth_file_from_env(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path / "env-dir"))
monkeypatch.setenv("CHATGPT_AUTH_FILE", "env.json")
authenticator: Final = Authenticator()
assert authenticator.auth_file == str(tmp_path / "env-dir" / "env.json")
def test_bare_filename_auth_file_resolves_against_cwd(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.chdir(tmp_path)
authenticator: Final = Authenticator(auth_file="auth.json")
assert authenticator.auth_file == str(tmp_path / "auth.json")
assert authenticator.token_dir == str(tmp_path)
def test_two_auth_files_are_isolated(self, tmp_path: Path) -> None:
file_a: Final = tmp_path / "account-a.json"
file_b: Final = tmp_path / "account-b.json"
_write_auth_record(file_a, "token-a", "acct-a")
_write_auth_record(file_b, "token-b", "acct-b")
authenticator_a: Final = Authenticator(auth_file=str(file_a))
authenticator_b: Final = Authenticator(auth_file=str(file_b))
assert authenticator_a.get_access_token() == "token-a"
assert authenticator_a.get_account_id() == "acct-a"
assert authenticator_b.get_access_token() == "token-b"
assert authenticator_b.get_account_id() == "acct-b"
def test_get_cached_authenticator_reuses_instance_per_path(self, tmp_path: Path) -> None:
file_a: Final = tmp_path / "account-a.json"
file_b: Final = tmp_path / "account-b.json"
authenticator_a: Final = get_cached_authenticator(str(file_a))
assert get_cached_authenticator(str(file_a)) is authenticator_a
assert get_cached_authenticator(str(file_b)) is not authenticator_a
def test_get_chatgpt_auth_file(self) -> None:
assert get_chatgpt_auth_file(None) is None
assert get_chatgpt_auth_file({}) is None
assert get_chatgpt_auth_file({"chatgpt_auth_file": ""}) is None
assert get_chatgpt_auth_file({"chatgpt_auth_file": "/a/auth.json"}) == "/a/auth.json"
assert get_chatgpt_auth_file(GenericLiteLLMParams()) is None
assert get_chatgpt_auth_file(GenericLiteLLMParams(chatgpt_auth_file="/b/auth.json")) == "/b/auth.json"
def test_relative_and_absolute_auth_file_share_cached_authenticator(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.chdir(tmp_path)
get_cached_authenticator.cache_clear()
auth_file: Final = tmp_path / "x" / "auth.json"
auth_file.parent.mkdir(parents=True)
auth_file.write_text("{}")
relative_key: Final = get_chatgpt_auth_file({"chatgpt_auth_file": "x/auth.json"})
absolute_key: Final = get_chatgpt_auth_file({"chatgpt_auth_file": str(auth_file)})
assert relative_key is not None and absolute_key is not None
assert get_cached_authenticator(relative_key) is get_cached_authenticator(absolute_key)
def test_refresh_keeps_saved_workspace(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr("time.time", lambda: 1700000000)
auth_file: Final = tmp_path / "selected.json"
auth_file.write_text(
json.dumps(
{
"access_token": "expired-example-token",
"refresh_token": "example-refresh-token",
"expires_at": 1,
"account_id": "acct-selected-workspace",
}
)
)
token: Final = _make_jwt({"exp": 4102444800, "https://api.openai.com/auth": {"chatgpt_account_id": "acct-default"}})
authenticator: Final = Authenticator(auth_file=str(auth_file))
with respx.mock as transport:
transport.post("https://auth.openai.com/oauth/token").mock(
return_value=httpx.Response(
200,
json={"access_token": token, "id_token": token, "refresh_token": "example-rotated-refresh"},
)
)
assert authenticator.get_access_token() == token
assert authenticator.get_account_id(token) == "acct-selected-workspace"
assert json.loads(auth_file.read_text())["refresh_token"] == "example-rotated-refresh"

View file

@ -0,0 +1,90 @@
import base64
import json
import pytest
from litellm.llms.chatgpt.common_utils import (
extract_chatgpt_account_id,
get_chatgpt_client_credential,
is_chatgpt_oauth_key,
select_chatgpt_client_credential_headers,
)
def _make_jwt(payload: dict) -> str:
def _b64(obj: dict) -> str:
raw = json.dumps(obj, separators=(",", ":")).encode("utf-8")
return base64.urlsafe_b64encode(raw).decode("utf-8").rstrip("=")
return f"{_b64({'alg': 'none', 'typ': 'JWT'})}.{_b64(payload)}."
CHATGPT_JWT = _make_jwt({"https://api.openai.com/auth": {"chatgpt_account_id": "acct-alice"}})
CHATGPT_JWT_WITHOUT_ACCOUNT = _make_jwt({"https://api.openai.com/auth": {"user_id": "user-1"}})
NON_CHATGPT_JWT = _make_jwt({"sub": "someone", "https://api.openai.com/auth": "not-an-object"})
class TestExtractChatGPTAccountId:
def test_reads_account_id_from_auth_claim(self):
assert extract_chatgpt_account_id(CHATGPT_JWT) == "acct-alice"
@pytest.mark.parametrize(
"token",
[None, "", "sk-proj-not-a-jwt", "header.only", NON_CHATGPT_JWT, CHATGPT_JWT_WITHOUT_ACCOUNT],
)
def test_returns_none_without_a_usable_account_claim(self, token):
assert extract_chatgpt_account_id(token) is None
class TestIsChatGPTOAuthKey:
def test_recognizes_bearer_and_raw_chatgpt_jwt(self):
assert is_chatgpt_oauth_key(CHATGPT_JWT)
assert is_chatgpt_oauth_key(f"Bearer {CHATGPT_JWT}")
assert is_chatgpt_oauth_key(f"Bearer {CHATGPT_JWT_WITHOUT_ACCOUNT}")
@pytest.mark.parametrize(
"value", [None, "", "Bearer sk-proj-platform-key", "Bearer not.a.jwt", _make_jwt({"sub": "x"})]
)
def test_rejects_non_chatgpt_credentials(self, value):
assert not is_chatgpt_oauth_key(value)
class TestSelectChatGPTClientCredentialHeaders:
def test_account_header_alone_is_not_forwarded(self):
selected = select_chatgpt_client_credential_headers(
{"chatgpt-account-id": "acct-alice", "authorization": "Bearer sk-proj-platform-key"}
)
assert dict(selected) == {}
assert get_chatgpt_client_credential({"chatgpt-account-id": "acct-alice"}) is None
def test_forwards_only_bearer_and_account_context(self):
selected = select_chatgpt_client_credential_headers(
{
"Authorization": f"Bearer {CHATGPT_JWT}",
"ChatGPT-Account-Id": "acct-explicit",
"x-litellm-api-key": "sk-gateway",
"user-agent": "codex",
}
)
assert dict(selected) == {"Authorization": f"Bearer {CHATGPT_JWT}", "ChatGPT-Account-Id": "acct-explicit"}
credential = get_chatgpt_client_credential(selected)
assert credential is not None
assert credential.access_token == CHATGPT_JWT
assert credential.account_id == "acct-explicit"
def test_account_id_falls_back_to_jwt_claim(self):
credential = get_chatgpt_client_credential({"authorization": f"Bearer {CHATGPT_JWT}"})
assert credential is not None
assert credential.account_id == "acct-alice"
@pytest.mark.parametrize("scheme", ["Bearer", "bearer", "BEARER"])
def test_bearer_scheme_is_case_insensitive(scheme):
credential = get_chatgpt_client_credential({"authorization": f"{scheme} {CHATGPT_JWT}"})
assert credential is not None
assert credential.access_token == CHATGPT_JWT
assert credential.account_id == "acct-alice"
def test_other_authentication_schemes_are_not_chatgpt_logins():
assert not is_chatgpt_oauth_key(f"Basic {CHATGPT_JWT}")

View file

@ -5654,3 +5654,44 @@ def test_vector_store_retrieve_handler_sends_id_before_azure_api_version_query()
sent: Final = received[0]
assert str(sent.url) == "https://x.openai.azure.com/openai/vector_stores/vs_abc?api-version=2025-04-01-preview"
assert len(received) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("is_async", (False, True))
@pytest.mark.parametrize("provider", ("openai", "azure"))
async def test_non_chatgpt_responses_preserves_caller_headers_after_validation(
is_async: bool, provider: str, respx_mock: respx.MockRouter,
) -> None:
route: Final = respx_mock.post(url__regex=r"https://responses\.example\.com/.*").respond(
200,
json={
"id": "resp_headers", "object": "response", "created_at": 1,
"model": "gpt-6-luna", "status": "completed", "output": [],
"parallel_tool_calls": False, "tool_choice": "auto", "tools": [],
},
)
client: Final = AsyncHTTPHandler(transport=httpx.AsyncHTTPTransport()) if is_async else HTTPHandler(client=httpx.Client())
kwargs: Final = {
"model": f"{provider}/gpt-6-luna",
"input": "hello",
"api_key": "fake-configured-key",
"api_base": "https://responses.example.com/v1",
"api_version": "2025-04-01-preview",
"extra_headers": {"Authorization": "Bearer fake-caller-key", "content-type": "application/vnd.example+json"},
"client": client,
}
try:
response: Final = await litellm.aresponses(**kwargs) if is_async else litellm.responses(**kwargs)
assert response.status == "completed"
assert route.calls[0].request.headers["authorization"] == "Bearer fake-caller-key"
expected_content_types: Final = (
["application/json", "application/vnd.example+json"]
if provider == "openai"
else ["application/vnd.example+json"]
)
assert route.calls[0].request.headers.get_list("content-type") == expected_content_types
finally:
if isinstance(client, AsyncHTTPHandler):
await client.close()
else:
client.client.close()

View file

@ -314,4 +314,4 @@ class TestElevenLabsTextToSpeechTransformation:
)
assert voice_id in url
assert "output_format=pcm_44100" in url
assert "output_format=pcm_24000" in url

View file

@ -55,3 +55,31 @@ def test_speech_keeps_an_internal_prefixed_kwarg_out_of_the_elevenlabs_request(r
sent: Final = json.loads(respx_mock.calls[0].request.content)
assert "_litellm_undeclared_sentinel" not in sent, sent
assert sent["text"] == "hi"
@pytest.mark.parametrize(
("response_format", "expected_output_format"),
[("wav", "wav_24000"), ("pcm", "pcm_24000")],
)
def test_speech_maps_wav_and_pcm_to_elevenlabs_output_formats(
respx_mock: respx.MockRouter,
response_format: str,
expected_output_format: str,
) -> None:
api_base: Final = "http://localhost:12346"
mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/v1/text-to-speech/.*").mock(
return_value=httpx.Response(status_code=200, content=b"audio", headers={"content-type": "audio/mpeg"})
)
litellm.speech(
model="elevenlabs/eleven_multilingual_v2",
input="hi",
voice="21m00Tcm4TlvDq8ikWAM",
response_format=response_format,
api_base=api_base,
api_key="fake_elevenlabs_api_key",
)
assert mock_route.called
request_url: Final = httpx.URL(str(respx_mock.calls[0].request.url))
assert request_url.params["output_format"] == expected_output_format

View file

@ -2,7 +2,7 @@ import base64
from collections.abc import Sequence
from itertools import chain
from pathlib import Path
from typing import Final
from typing import Final, Literal
import httpx
import pytest
@ -1650,10 +1650,28 @@ def test_convert_tool_response_with_base64_image():
assert inline_data["data"] == test_image_base64
def test_gemini_history_nests_multimodal_tool_response_parts():
"""Full history conversion should not emit sibling inline_data tool result parts."""
test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
messages = [
@pytest.mark.parametrize("provider", ["gemini", "vertex_ai", "vertex_ai_beta"])
@pytest.mark.parametrize(
"model,nested",
[
(None, True),
("gemini-2.5-pro", False),
("gemini/gemini-2.5-flash", False),
("vertex_ai/gemini-2.5-flash-lite", False),
("123456789", False),
("gemini-3.8-flash", True),
("gemini/gemini-3.8-flash", True),
("vertex_ai/gemini-3.8-flash", True),
("gemini-flash-latest", True),
],
)
def test_gemini_history_places_tool_media_for_model(
model: str | None,
nested: bool,
provider: Literal["gemini", "vertex_ai", "vertex_ai_beta"],
) -> None:
test_image_base64: Final = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
messages: Final = [
{"role": "user", "content": "Get me an image"},
{
"role": "assistant",
@ -1679,24 +1697,53 @@ def test_gemini_history_nests_multimodal_tool_response_parts():
"data": test_image_base64,
},
},
{
"type": "file",
"file": {"file_data": "data:application/pdf;base64,JVBERi0xLjQ="},
},
],
},
{"role": "user", "content": "Describe the image and file again"},
]
contents = gemini_convert_messages_with_history(messages=messages)
tool_response_parts = contents[-1]["parts"]
assert len(tool_response_parts) == 1
assert "inline_data" not in tool_response_parts[0]
function_response = tool_response_parts[0]["function_response"]
assert function_response["parts"] == [
contents: Final = (
gemini_convert_messages_with_history(messages=messages)
if model is None
else transform_request_body(
messages=messages,
model=model,
custom_llm_provider=provider,
optional_params={},
litellm_params={},
cached_content=None,
)["contents"]
)
media_parts: Final = [
{
"inline_data": {
"data": test_image_base64,
"mime_type": "image/png",
}
}
},
{"inline_data": {"data": "JVBERi0xLjQ=", "mime_type": "application/pdf"}},
]
function_response: Final = {
"name": "get_image",
"response": {"image_ref": "inline"},
**({"id": "call_get_image"} if model is not None and nested else {}),
**({"parts": media_parts} if nested else {}),
}
assert contents[-2] == {
"role": "user",
"parts": [
{"function_response": function_response},
*([] if nested else media_parts),
],
}
assert contents[-1] == {
"role": "user",
"parts": [{"text": "Describe the image and file again"}],
}
def test_convert_tool_response_text_only():

View file

@ -7693,6 +7693,47 @@ class TestMCPDcrBridgeDualCredential:
assert "invalid_token" in str(exc_info.value.headers)
@pytest.mark.parametrize("public_path", ["/toolset/analysis/mcp", "/analysis/mcp", "/mcp", "/mcp/", "/mcp/proxy"])
@pytest.mark.parametrize("root_path", ["", "/litellm"])
@pytest.mark.parametrize("invalid_token", [False, True])
def test_gateway_challenge_metadata_contains_the_public_resource(
public_path: str, root_path: str, invalid_token: bool,
) -> None:
from mcp.shared.auth_utils import check_resource_allowed
from starlette.requests import Request
from starlette.types import Scope
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import _gateway_dcr_challenge
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_build_aggregate_protected_resource_response,
oauth_protected_resource_root,
)
scope: Final[Scope] = {
"type": "http", "method": "POST", "scheme": "https", "path": "/mcp",
"root_path": root_path, "_original_path": f"{root_path}{public_path}",
"headers": [(b"host", b"gateway.example.com")],
}
request: Final = Request(scope)
challenge: Final = _gateway_dcr_challenge(request, "/mcp", None, invalid_token)
header: Final = (challenge.headers or {})["WWW-Authenticate"]
metadata_url: Final = header.split('resource_metadata="', 1)[1].split('"', 1)[0]
toolset: Final = public_path in ("/toolset/analysis/mcp", "/analysis/mcp")
assert metadata_url == (
f"https://gateway.example.com{root_path}/.well-known/oauth-protected-resource"
+ ("" if toolset else "/mcp")
)
metadata: Final = (
oauth_protected_resource_root(request) if toolset else _build_aggregate_protected_resource_response(request)
)
assert check_resource_allowed(
requested_resource=f"https://gateway.example.com{root_path}{public_path}",
configured_resource=metadata["resource"],
)
assert challenge.status_code == 401
assert ('error="invalid_token"' in header) is invalid_token
@pytest.mark.asyncio
class TestAggregateGatewayDcrChallenge:
"""The mcp_gateway_dcr front door: a 401 on the aggregate /mcp scope must

View file

@ -5,6 +5,7 @@ from litellm.proxy._experimental.mcp_server import operations as mcp_operations
import asyncio
import logging
import sys
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import httpx
@ -531,12 +532,12 @@ async def test_scope_challenge_replaces_pending_stream_headers_once():
send = AsyncMock()
response = OAuthScopeResponse(send, "https://gateway.example", "/docs/mcp")
await response({"type": "http.response.start", "status": 200, "headers": []})
headers: Final = asyncio.create_task(response({"type": "http.response.start", "status": 200, "headers": []}))
body = asyncio.create_task(response({"type": "http.response.body", "body": b"event: message\n"}))
await asyncio.sleep(0)
send.assert_not_awaited()
response.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write"))
await asyncio.wait_for(body, 1)
await asyncio.wait_for(asyncio.gather(headers, body), 1)
await response({"type": "http.response.body", "body": b"private tool error"})
assert send.await_count == 2
start, result = [call.args[0] for call in send.await_args_list]
@ -567,7 +568,7 @@ async def test_scope_authorization_releases_progress_without_waiting_for_tool_re
context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=request)
token = active_mcp_request_ctx_var.set(context)
try:
await gate({"type": "http.response.start", "status": 200, "headers": []})
headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []}))
progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True}))
await asyncio.sleep(0)
for method, status in [("initialize", 200), ("tools/list", 200), ("tools/call", 403)]:
@ -582,7 +583,7 @@ async def test_scope_authorization_releases_progress_without_waiting_for_tool_re
upstream = MCPClient(server_url="https://upstream.example/mcp")
async with upstream._create_httpx_client_factory(transport=httpx2.MockTransport(lambda request: httpx2.Response(200, content=b"tool result")))() as client:
async with client.stream("POST", "https://upstream.example/mcp", json={"method": "tools/call"}):
await asyncio.wait_for(progress, 1)
await asyncio.wait_for(asyncio.gather(headers, progress), 1)
assert send.await_args_list[-1].args[0]["body"] == b"progress"
gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write"))
await gate({"type": "http.response.body", "body": b"result"})
@ -594,21 +595,105 @@ async def test_scope_authorization_releases_progress_without_waiting_for_tool_re
@pytest.mark.asyncio
@pytest.mark.parametrize("failure", [ValueError, asyncio.CancelledError])
async def test_scope_middleware_releases_waiting_response_on_validation_or_cancellation(failure):
@pytest.mark.parametrize("fail_send", [False, True])
async def test_scope_middleware_releases_waiting_response_on_validation_or_cancellation(
failure: type[BaseException], fail_send: bool,
) -> None:
from mcp.server.context import ServerRequestContext
from starlette.requests import Request
from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse, finish_scope_response
send = AsyncMock()
gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs")
context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate}))
send: Final = AsyncMock(side_effect=BrokenPipeError("peer disconnected") if fail_send else None)
gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs")
context: Final = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate}))
headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []}))
await asyncio.sleep(0)
with pytest.raises(failure):
await finish_scope_response(context, AsyncMock(side_effect=failure))
await gate({"type": "http.response.start", "status": 200, "headers": []})
if fail_send:
with pytest.raises(BrokenPipeError, match="peer disconnected"):
await headers
else:
await headers
await gate({"type": "http.response.body", "body": b"error"})
assert not gate.pending
assert send.await_args_list[-1].args[0]["body"] == b"error"
assert send.await_args_list[-1].args[0]["type"] == ("http.response.start" if fail_send else "http.response.body")
@pytest.mark.asyncio
@pytest.mark.parametrize("fail_send", [False, True])
async def test_scope_transport_owns_header_send_cancellation_and_errors(fail_send: bool) -> None:
import httpx2
from starlette.types import Message
from litellm.proxy._experimental.mcp_server.legacy_callbacks import record_upstream_tool_authorization
from litellm.proxy._experimental.mcp_server.scope_challenge import OAuthScopeResponse
send_started: Final = asyncio.Event()
send_stopped: Final = asyncio.Event()
writable: Final = asyncio.Event()
sent: Final[list[Message]] = []
async def send(message: Message) -> None:
sent.append(message)
assert message["type"] == "http.response.start"
send_started.set()
try:
if fail_send:
raise BrokenPipeError("peer disconnected")
await writable.wait()
finally:
send_stopped.set()
gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs")
headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []}))
body: Final = asyncio.create_task(gate({"type": "http.response.body", "body": b"result"}))
await asyncio.sleep(0)
assert not send_started.is_set()
response: Final = httpx2.Response(200, request=httpx2.Request(
"POST", "https://upstream.example/mcp", json={"method": "tools/call"},
))
try:
await record_upstream_tool_authorization(response, pending=gate)
await send_started.wait()
if fail_send:
with pytest.raises(BrokenPipeError, match="peer disconnected"):
await headers
else:
headers.cancel()
with pytest.raises(asyncio.CancelledError):
await headers
await body
assert len(sent) == 1
assert send_stopped.is_set()
finally:
writable.set()
await send_stopped.wait()
@pytest.mark.asyncio
@pytest.mark.parametrize("allowed_before_start", [False, True])
async def test_scope_authorization_sends_stream_headers_before_the_first_body_chunk(allowed_before_start: bool) -> None:
from starlette.types import Message
from litellm.proxy._experimental.mcp_server.scope_challenge import OAuthScopeResponse
send: Final = AsyncMock()
gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs")
start: Final[Message] = {"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/event-stream")]}
if allowed_before_start:
gate.allow()
headers: Final = asyncio.create_task(gate(start))
await asyncio.sleep(0)
if not allowed_before_start:
send.assert_not_awaited()
gate.allow()
await headers
send.assert_awaited_once_with(start)
await gate({"type": "http.response.body", "body": b"event: message\n", "more_body": True})
assert send.await_args_list[-1].args[0]["body"] == b"event: message\n"
assert send.await_count == 2
@pytest.mark.asyncio
@ -697,13 +782,14 @@ async def test_scope_challenges_are_isolated_between_concurrent_requests():
context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate}))
token = active_mcp_request_ctx_var.set(context)
try:
await gate({"type": "http.response.start", "status": 200, "headers": []})
headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []}))
await asyncio.sleep(0)
assert get_scope_response() is gate
if denied:
gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write"))
else:
await record_upstream_tool_authorization(httpx2.Response(200, request=httpx2.Request("POST", "https://upstream.example/mcp", json={"method": "tools/call"})), pending=get_scope_response())
await headers
await gate({"type": "http.response.body", "body": b"done"})
return send.await_args_list[0].args[0]["status"]
finally:

View file

@ -7681,7 +7681,7 @@ class TestMCPServerManager:
request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate}),
)
token = active_mcp_request_ctx_var.set(context)
await gate({"type": "http.response.start", "status": 200, "headers": []})
headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []}))
progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True}))
async def upstream_call(*args, **kwargs):
@ -7701,7 +7701,7 @@ class TestMCPServerManager:
assert not result.is_error
finally:
gate.allow()
await asyncio.wait_for(progress, 1)
await asyncio.wait_for(asyncio.gather(headers, progress), 1)
active_mcp_request_ctx_var.reset(token)
@pytest.mark.asyncio

View file

@ -280,7 +280,7 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(
send = AsyncMock()
response = OAuthScopeResponse(send, "https://gateway.example", "/mcp/pt")
request = Request({"type": "http", "headers": [], SCOPE_RESPONSE_KEY: response})
await response({"type": "http.response.start", "status": 200, "headers": []})
headers: Final = asyncio.create_task(response({"type": "http.response.start", "status": 200, "headers": []}))
try:
from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError
from litellm.proxy._experimental.mcp_server.server import (
@ -328,6 +328,7 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(
info_calls = [str(c.args[0]) for c in mock_logger.info.call_args_list if c.args]
assert any("Upstream auth failure" in m for m in info_calls)
response.allow()
await headers
await response({"type": "http.response.body", "body": b"tool error"})
start = send.await_args_list[0].args[0]
assert start["status"] == (403 if required_scope is not None else 200)

View file

@ -3626,6 +3626,18 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields:
)
@pytest.mark.parametrize("allow_client_side_credentials", [False, True])
@pytest.mark.parametrize("nested", [False, True])
def test_request_cannot_select_chatgpt_auth_file(allow_client_side_credentials: bool, nested: bool) -> None:
selector: Final = {"chatgpt_auth_file": "/tokens/account-b/auth.json"}
with pytest.raises(ValueError, match="chatgpt_auth_file"):
is_request_body_safe(
request_body={"model": "chatgpt-account-a", **({"extra_body": selector} if nested else selector)},
general_settings={"allow_client_side_credentials": allow_client_side_credentials},
llm_router=None,
model="chatgpt-account-a",
)
class TestIsRequestBodySafeBlocksBedrockProjectOverride:
"""``aws_bedrock_project_id`` pins a deployment to a Bedrock project so
that project's data-retention policy applies to its requests. A
@ -5304,3 +5316,27 @@ async def test_a_non_numeric_max_request_size_mb_is_ignored_without_premium():
patch("litellm.proxy.proxy_server.premium_user", False),
):
assert await check_if_request_size_is_safe(request=_request_with_content_length(1024)) is True
@pytest.mark.parametrize("allow_client_side_credentials", [False, True])
@pytest.mark.parametrize("route", ["/model/new", "/model/update", "/model/account-a/update", "/health/test_connection"])
def test_chatgpt_model_configuration_reaches_role_authorization(
allow_client_side_credentials: bool, route: str
) -> None:
body: Final = {"litellm_params": {"model": "chatgpt/gpt-6-luna", "chatgpt_auth_file": "/tokens/account-a/auth.json"}}
assert is_request_body_safe(
request_body=body,
general_settings={"allow_client_side_credentials": allow_client_side_credentials},
llm_router=None,
model="",
route=route,
)
for inference_route in ("/v1/chat/completions", "/v1/responses"):
with pytest.raises(ValueError, match="chatgpt_auth_file"):
is_request_body_safe(
request_body=body,
general_settings={"allow_client_side_credentials": allow_client_side_credentials},
llm_router=None,
model="",
route=inference_route,
)

View file

@ -1,8 +1,9 @@
from collections.abc import Mapping
from typing import Final, TypedDict
from typing import Final, TypedDict, cast
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import HTTPException
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
from litellm.proxy import proxy_server
@ -78,6 +79,14 @@ def _team_loader(team_tools: list[str] | None | str) -> TeamObjectLoader:
return load
def _missing_dashboard_team_loader() -> tuple[TeamObjectLoader, AsyncMock]:
missing_team: Final = HTTPException(
status_code=404, detail="Team doesn't exist in db. Team=litellm-dashboard"
)
load_team: Final = AsyncMock(side_effect=missing_team)
return cast(TeamObjectLoader, load_team), load_team
async def _no_team() -> None:
return None
@ -194,9 +203,6 @@ async def test_search_tool_access_is_unchanged_without_search_tool_deny_by_defau
id="proxy admin virtual key not exempt",
),
pytest.param(CallerFields(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS), None, None, None, id="master key exempt"),
pytest.param(
CallerFields(team_id=UI_TEAM_ID, virtual_key=False), None, None, None, id="dashboard session exempt"
),
],
)
@pytest.mark.asyncio
@ -248,6 +254,79 @@ async def test_search_tool_deny_by_default_does_not_load_the_user_when_off(
get_user_object.assert_not_awaited()
@pytest.mark.parametrize(
"general_settings",
[{}, {"search_tool_deny_by_default": False}],
ids=["omitted", "false"],
)
@pytest.mark.parametrize(
"user_role",
[LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.INTERNAL_USER],
ids=["admin", "internal-user"],
)
@pytest.mark.parametrize(
"key_tools, expected",
[
pytest.param(None, None, id="unrestricted-key"),
pytest.param(["search-a"], None, id="key-grants-search-a"),
pytest.param(["search-b"], _LEGACY, id="key-excludes-search-a"),
],
)
@pytest.mark.asyncio
async def test_dashboard_search_tool_grants_without_deny_by_default(
general_settings: Mapping[str, object],
user_role: LitellmUserRoles,
key_tools: list[str] | None,
expected: ProxyErrorTypes | None,
) -> None:
caller: Final = _caller(
virtual_key=False,
key_tools=key_tools,
team_id=UI_TEAM_ID,
user_role=user_role,
)
load_team, load_team_mock = _missing_dashboard_team_loader()
assert await _denied_by(general_settings, caller, load_team) == expected
load_team_mock.assert_not_awaited()
@pytest.mark.parametrize(
"user_role",
[LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY],
ids=["admin", "admin-view-only"],
)
@pytest.mark.asyncio
async def test_admin_dashboard_search_tool_grants_with_deny_by_default(user_role: LitellmUserRoles) -> None:
caller: Final = _caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=user_role)
load_team, load_team_mock = _missing_dashboard_team_loader()
assert await _denied_by(_DENY_ON, caller, load_team) is None
load_team_mock.assert_not_awaited()
@pytest.mark.parametrize(
"user_id, user_tools, expected",
[
pytest.param("user-1", ["search-a"], None, id="user-grants-search-a"),
pytest.param("user-1", ["search-b"], _USER, id="user-grants-another-tool"),
pytest.param("user-1", None, _USER, id="user-grant-null"),
pytest.param("user-2", None, _USER, id="no-user-row"),
],
)
@pytest.mark.asyncio
async def test_internal_dashboard_search_tool_requires_the_user_grant_with_deny_by_default(
cache: UserApiKeyCache,
user_id: str,
user_tools: list[str] | None,
expected: ProxyErrorTypes | None,
) -> None:
if user_id == "user-1":
_cache_user(cache, user_tools)
caller: Final = _caller(virtual_key=False, team_id=UI_TEAM_ID, user_id=user_id)
load_team, load_team_mock = _missing_dashboard_team_loader()
assert await _denied_by(_DENY_ON, caller, load_team) == expected
load_team_mock.assert_not_awaited()
@pytest.mark.parametrize(
"general_settings, caller, expected",
[
@ -257,6 +336,24 @@ async def test_search_tool_deny_by_default_does_not_load_the_user_when_off(
pytest.param(_DENY_ON, _caller(virtual_key=False), _USER, id="keyless user"),
pytest.param(_DENY_ON, _caller(user_role=LitellmUserRoles.PROXY_ADMIN), _KEY, id="proxy admin key"),
pytest.param(_DENY_ON, _caller(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS), None, id="master key"),
pytest.param(
_DENY_ON,
_caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=LitellmUserRoles.INTERNAL_USER),
_USER,
id="internal dashboard user",
),
pytest.param(
_DENY_ON,
_caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=LitellmUserRoles.PROXY_ADMIN),
None,
id="admin dashboard",
),
pytest.param(
_DENY_ON,
_caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY),
None,
id="admin view-only dashboard",
),
],
)
def test_unregistered_search_fallback_follows_search_tool_deny_by_default(
@ -268,3 +365,12 @@ def test_unregistered_search_fallback_follows_search_tool_deny_by_default(
with pytest.raises(ProxyException) as exc_info:
check_unregistered_search_fallback(caller, general_settings)
assert (exc_info.value.type, exc_info.value.code) == (expected, "403")
@pytest.mark.asyncio
async def test_search_preserves_missing_real_team_error_without_deny_by_default() -> None:
missing_team: Final = HTTPException(status_code=404, detail="Team doesn't exist in db. Team=team-1")
load_team: Final = AsyncMock(side_effect=missing_team)
with pytest.raises(HTTPException) as exc_info:
await can_caller_call_search_tool("search-a", _caller(team_id="team-1"), {}, load_team)
assert exc_info.value is missing_team

View file

@ -5437,3 +5437,64 @@ async def test_health_services_endpoint_langfuse_missing_keys_errors(monkeypatch
with pytest.raises(ProxyException, match="auth_check failed") as raised:
await health_services_endpoint(service="langfuse")
assert "connection refused by lf.internal.example" in str(raised.value.message)
@pytest.mark.asyncio
@pytest.mark.parametrize("admin", [False, True])
async def test_unsaved_chatgpt_account_probe_requires_proxy_admin(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, admin: bool
) -> None:
from starlette.requests import Request
from litellm.proxy import proxy_server
monkeypatch.setenv("CHATGPT_AUTH_FILE", str(tmp_path / "missing-default-auth.json"))
auth_file: Final = tmp_path / "auth.json"
auth_file.write_text(
json.dumps({"access_token": "token-probe", "account_id": "acct-probe", "expires_at": 4102444800})
)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
monkeypatch.setattr(proxy_server, "llm_router", None)
monkeypatch.setattr(proxy_server, "general_settings", {})
output: Final = {
"id": "msg_probe",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "ok", "annotations": []}],
}
response: Final = {
"id": "resp_probe",
"object": "response",
"created_at": 1700000000,
"status": "completed",
"model": "gpt-6.1-sol",
"output": [output],
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
}
with respx.mock(assert_all_called=False) as upstream:
upstream.post("https://chatgpt.com/backend-api/codex/responses").respond(
200,
text="data: " + json.dumps({"type": "response.completed", "response": response}) + "\n\n",
headers={"content-type": "text/event-stream"},
)
call: Final = health_test_model_connection(
request=Request({"type": "http", "headers": []}),
mode="responses",
litellm_params={"model": "chatgpt/gpt-6.1-sol", "chatgpt_auth_file": str(auth_file)},
model_info={},
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN if admin else LitellmUserRoles.INTERNAL_USER, token="test-key"
),
)
if admin:
result: Final = await call
assert result["status"] == "success"
assert upstream.calls.last.request.headers["Authorization"] == "Bearer token-probe"
assert upstream.calls.last.request.headers["ChatGPT-Account-Id"] == "acct-probe"
assert isinstance(json.loads(upstream.calls.last.request.content)["input"], list)
else:
with pytest.raises(ProxyException) as error:
await call
assert error.value.code == "403"
assert len(upstream.calls) == 0

View file

@ -7,7 +7,7 @@ import contextvars
import logging
import sys
import time
from collections.abc import Iterator, Sequence
from collections.abc import AsyncGenerator, Iterator, Sequence
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Final, List, Optional
@ -4644,6 +4644,140 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch):
)
_SSE_KEEPALIVE_CALLBACK_TIME: Final = datetime(2026, 1, 1)
@pytest.mark.asyncio
@pytest.mark.parametrize("after_ping", [False, True], ids=["fast", "pinged"])
@pytest.mark.parametrize("fails", [False, True], ids=["success", "failure"])
@pytest.mark.parametrize("preadmitted", [False, True], ids=["producer-admission", "existing-admission"])
async def test_sse_keepalive_releases_parallel_capacity(
after_ping: bool, fails: bool, preadmitted: bool
) -> None:
from fastapi.responses import StreamingResponse
from litellm.proxy.common_request_processing import open_sse_before_first_byte
from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_BYTES
cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
api_key: Final = hash_token("sk-keepalive-release")
auth: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
counter: Final = f"{{api_key:{api_key}}}:max_parallel_requests"
async def request(call_id: str) -> None:
data: Final = {"model": "test-model", "litellm_call_id": call_id}
upstream_answers: Final = asyncio.Event()
if not after_ping:
upstream_answers.set()
async def admit() -> None:
await handler.async_pre_call_hook(
user_api_key_dict=auth, cache=cache, data=data, call_type="acompletion"
)
async def failure(exc: Exception) -> None:
await handler.async_post_call_failure_hook(
request_data=data, original_exception=exc, user_api_key_dict=auth
)
async def stream() -> AsyncGenerator[bytes, None]:
yield b"data: [DONE]\n\n"
await handler.async_log_success_event(
kwargs={
"litellm_call_id": call_id,
"standard_logging_object": {"metadata": {"user_api_key_hash": api_key}},
},
response_obj=ModelResponse(usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2)),
start_time=_SSE_KEEPALIVE_CALLBACK_TIME,
end_time=_SSE_KEEPALIVE_CALLBACK_TIME,
)
async def produce() -> StreamingResponse:
if not preadmitted:
await admit()
await upstream_answers.wait()
if fails:
raise HTTPException(status_code=502, detail="upstream failed")
return StreamingResponse(stream())
if preadmitted:
await admit()
response_call: Final = open_sse_before_first_byte(
produce(), 0.001 if after_ping else 60.0, on_late_failure=failure
)
if fails and not after_ping:
with pytest.raises(HTTPException) as rejected:
await response_call
assert rejected.value.status_code == 502
await failure(rejected.value)
return
response: Final = await response_call
first: Final = await anext(response.body_iterator)
if after_ping:
upstream_answers.set()
rest: Final = tuple([chunk async for chunk in response.body_iterator])
assert first == (SSE_COMMENT_PING_BYTES if after_ping else b"data: [DONE]\n\n")
if fails:
error_frame: Final = rest[-2]
assert isinstance(error_frame, bytes) and b"upstream failed" in error_frame
assert rest[-1] == b"data: [DONE]\n\n"
elif after_ping:
assert rest[-1] == b"data: [DONE]\n\n"
else:
assert rest == ()
for index in range(3):
await contextvars.Context().run(asyncio.create_task, request(f"keepalive-{index}"))
assert handler._gauge_in_flight_from_cache_value(await cache.async_get_cache(counter)) == 0
@pytest.mark.asyncio
async def test_sse_keepalive_rejects_overlap_and_reuses_completed_capacity() -> None:
from fastapi.responses import StreamingResponse
from litellm.proxy.common_request_processing import open_sse_before_first_byte
cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
api_key: Final = hash_token("sk-keepalive-overlap")
auth: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
admitted: Final = asyncio.Event()
finish: Final = asyncio.Event()
async def request(call_id: str) -> None:
async def stream() -> AsyncGenerator[bytes, None]:
await finish.wait()
yield b"data: [DONE]\n\n"
await handler.async_post_call_success_hook(
data={"litellm_call_id": call_id}, user_api_key_dict=auth, response=ModelResponse()
)
async def produce() -> StreamingResponse:
await handler.async_pre_call_hook(
user_api_key_dict=auth,
cache=cache,
data={"model": "test-model", "litellm_call_id": call_id},
call_type="acompletion",
)
admitted.set()
return StreamingResponse(stream())
response: Final = await open_sse_before_first_byte(produce(), 60.0)
assert tuple([chunk async for chunk in response.body_iterator]) == (b"data: [DONE]\n\n",)
first: Final = contextvars.Context().run(asyncio.create_task, request("first"))
await admitted.wait()
try:
with pytest.raises(HTTPException) as rejected:
await contextvars.Context().run(asyncio.create_task, request("overlapping"))
assert rejected.value.status_code == 429
finally:
finish.set()
await first
await contextvars.Context().run(asyncio.create_task, request("after-completion"))
@pytest.mark.asyncio
async def test_read_only_gauge_check_counts_without_acquiring_v3():
"""

View file

@ -8933,7 +8933,14 @@ class TestNonAdminCannotPersistWifFieldsOnModel:
mock_prisma.db.litellm_proxymodeltable.update.assert_not_called()
@pytest.mark.asyncio
async def test_patch_model_non_admin_cannot_set_wif_field(self):
@pytest.mark.parametrize(
"field,value",
[
("anthropic_keycloak_token_url", "https://identity.example.com/token"),
("chatgpt_auth_file", "/tokens/account-b/auth.json"),
],
)
async def test_patch_model_non_admin_cannot_set_wif_field(self, field: str, value: str) -> None:
from litellm.proxy.management_endpoints.model_management_endpoints import (
patch_model,
)
@ -8979,13 +8986,13 @@ class TestNonAdminCannotPersistWifFieldsOnModel:
model_id="m1",
patch_data=updateDeployment(
litellm_params=updateLiteLLMParams(
anthropic_keycloak_token_url="https://attacker.example/token",
**{field: value},
)
),
user_api_key_dict=non_admin,
)
err = exc_info.value
assert getattr(err, "param", "") == "anthropic_keycloak_token_url"
assert getattr(err, "param", "") == field
mock_prisma.db.litellm_proxymodeltable.update.assert_not_called()
@pytest.mark.asyncio
@ -9044,7 +9051,14 @@ class TestNonAdminCannotPersistWifFieldsOnModel:
mock_prisma.db.litellm_proxymodeltable.update.assert_not_called()
@pytest.mark.asyncio
async def test_patch_model_admin_can_set_wif_field(self):
@pytest.mark.parametrize(
"field,value",
[
("anthropic_keycloak_token_url", "https://identity.example.com/token"),
("chatgpt_auth_file", "/tokens/account-b/auth.json"),
],
)
async def test_patch_model_admin_can_set_wif_field(self, field: str, value: str) -> None:
from litellm.proxy.management_endpoints.model_management_endpoints import (
patch_model,
)
@ -9095,7 +9109,7 @@ class TestNonAdminCannotPersistWifFieldsOnModel:
model_id="m1",
patch_data=updateDeployment(
litellm_params=updateLiteLLMParams(
anthropic_keycloak_token_url="https://keycloak.internal/token",
**{field: value},
)
),
user_api_key_dict=admin,

View file

@ -0,0 +1,185 @@
import asyncio
import json
from collections.abc import Iterator
from datetime import datetime
from typing import Final
import httpx
import pytest
import respx
from fastapi.testclient import TestClient
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy import proxy_server
from tests._master_key import MASTER_KEY
_BASE: Final = "https://chatgpt.com/backend-api/codex"
_TOKEN: Final = "test-subscription-token"
_HEADERS: Final = {"authorization": f"Bearer {_TOKEN}", "x-litellm-api-key": MASTER_KEY}
@pytest.fixture
def client(monkeypatch: pytest.MonkeyPatch, httpx_transport: None) -> Iterator[TestClient]:
monkeypatch.setattr(proxy_server, "master_key", MASTER_KEY)
yield TestClient(proxy_server.app)
def test_native_catalog_preserves_account_and_filters_gateway_credentials(
client: TestClient, respx_mock: respx.MockRouter
) -> None:
catalog: Final = {"models": [{"slug": "subscription-model", "supported_reasoning_levels": ["high"]}]}
upstream: Final = respx_mock.get(f"{_BASE}/models", params={"client_version": "1.0.0"}).respond(json=catalog)
response: Final = client.get(
"/backend-api/codex/models?client_version=1.0.0",
headers={
**_HEADERS,
"chatgpt-account-id": "account-a",
"cookie": "session=test-cookie",
"x-pass-authorization": "Bearer test-injected-token",
"x-pass-x-litellm-api-key": "test-injected-key",
},
)
assert response.status_code == 200
assert response.json() == catalog
assert upstream.calls.last.request.headers["authorization"] == f"Bearer {_TOKEN}"
assert upstream.calls.last.request.headers["chatgpt-account-id"] == "account-a"
assert "x-litellm-api-key" not in upstream.calls.last.request.headers
assert "cookie" not in upstream.calls.last.request.headers
@pytest.mark.parametrize("account", ["account-a", "account-b"])
def test_native_responses_preserves_body_and_stream(
client: TestClient, respx_mock: respx.MockRouter, account: str
) -> None:
body: Final = {
"model": "subscription-model",
"input": [{"role": "user", "content": "Hello"}],
"reasoning": {"effort": "high"},
"store": False,
"stream": True,
}
stream: Final = (
b'event: response.output_text.delta\ndata: {"type":"response.output_text.delta","delta":"Hello"}\n\n'
)
upstream: Final = respx_mock.post(f"{_BASE}/responses").respond(
content=stream, headers={"content-type": "text/event-stream"}
)
response: Final = client.post(
"/backend-api/codex/responses",
json=body,
headers={**_HEADERS, "authorization": f"Bearer test-{account}", "chatgpt-account-id": account},
)
assert response.status_code == 200
assert response.content == stream
assert response.headers["content-type"].startswith("text/event-stream")
assert json.loads(upstream.calls.last.request.content) == body
assert upstream.calls.last.request.headers["authorization"] == f"Bearer test-{account}"
assert upstream.calls.last.request.headers["chatgpt-account-id"] == account
@pytest.mark.parametrize("gateway_key", [None, "sk-invalid-test-key"])
def test_gateway_auth_rejects_before_upstream(
client: TestClient, respx_mock: respx.MockRouter, gateway_key: str | None
) -> None:
headers: Final = {
"authorization": f"Bearer {_TOKEN}",
**({"x-litellm-api-key": gateway_key} if gateway_key else {}),
}
response: Final = client.get("/backend-api/codex/models", headers=headers)
assert response.status_code in (400, 401, 403)
assert not respx_mock.calls
@pytest.mark.parametrize("authorization", ["", "Basic test", "Bearer ", f"Bearer {MASTER_KEY}"])
def test_subscription_credential_required(client: TestClient, respx_mock: respx.MockRouter, authorization: str) -> None:
response: Final = client.get("/backend-api/codex/models", headers={**_HEADERS, "authorization": authorization})
assert response.status_code == 401
assert not respx_mock.calls
@pytest.mark.parametrize("status", [401, 403, 429])
def test_provider_errors_do_not_fall_back(client: TestClient, respx_mock: respx.MockRouter, status: int) -> None:
upstream: Final = respx_mock.post(f"{_BASE}/responses").respond(
status, json={"error": {"message": "subscription unavailable"}}, headers={"retry-after": "60"}
)
response: Final = client.post(
"/backend-api/codex/responses", headers=_HEADERS, json={"model": "subscription-model", "stream": True}
)
assert response.status_code == status
assert upstream.call_count == 1
@pytest.mark.parametrize("method,path", [("POST", "models"), ("GET", "responses"), ("GET", "deviceauth")])
def test_only_native_inference_routes_are_forwarded(
client: TestClient, respx_mock: respx.MockRouter, method: str, path: str
) -> None:
response: Final = client.request(method, f"/backend-api/codex/{path}", headers=_HEADERS)
assert response.status_code in (405, 422)
assert not respx_mock.calls
class UsageRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.events: Final[asyncio.Queue[dict[str, object]]] = asyncio.Queue()
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
await self.events.put(kwargs)
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True])
async def test_usage_logged_without_credentials(
monkeypatch: pytest.MonkeyPatch, httpx_transport: None, respx_mock: respx.MockRouter, stream: bool
) -> None:
monkeypatch.setattr(proxy_server, "master_key", MASTER_KEY)
recorder: Final = UsageRecorder()
monkeypatch.setattr(litellm, "_async_success_callback", [recorder])
litellm.register_model(
{
"chatgpt/subscription-model": {
"litellm_provider": "chatgpt",
"mode": "responses",
"input_cost_per_token": 0.0,
"output_cost_per_token": 0.0,
}
}
)
result: Final = {
"id": "resp_test",
"object": "response",
"created_at": 1700000000,
"status": "completed",
"model": "subscription-model",
"output": [],
"usage": {"input_tokens": 12, "output_tokens": 3, "total_tokens": 15},
}
event: Final = {"type": "response.completed", "response": result}
response_body: Final = (
httpx.Response(
200,
content=f"event: response.completed\ndata: {json.dumps(event)}\n\n",
headers={"content-type": "text/event-stream"},
)
if stream
else httpx.Response(200, json=result)
)
respx_mock.post(f"{_BASE}/responses").mock(return_value=response_body)
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test") as caller:
response: Final = await caller.post(
"/backend-api/codex/responses",
json={"model": "subscription-model", "stream": stream, "input": "Hello"},
headers={**_HEADERS, "x-pass-authorization": "Bearer test-hidden-token", "cookie": "test-hidden-cookie"},
)
assert response.status_code == 200
logged: Final = await asyncio.wait_for(recorder.events.get(), timeout=5)
payload: Final = logged["standard_logging_object"]
assert isinstance(payload, dict)
assert payload["prompt_tokens"] == 12
assert payload["completion_tokens"] == 3
serialized: Final = json.dumps(logged, default=str)
for secret in (_TOKEN, MASTER_KEY, "test-hidden-token", "test-hidden-cookie"):
assert secret not in serialized

View file

@ -1011,10 +1011,9 @@ class TestPollingConditionChecks:
# Create mock router
mock_router = Mock()
mock_router.model_name_to_deployment_indices = {"gpt-5": [0]}
mock_router.model_list = [
{"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5"}}
]
mock_router.deployments_named.side_effect = {
"gpt-5": ({"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5"}},)
}.get
result = should_use_polling_for_request(
background_mode=True,
@ -1033,13 +1032,14 @@ class TestPollingConditionChecks:
)
mock_router = Mock()
mock_router.model_name_to_deployment_indices = {"claude-3": [0]}
mock_router.model_list = [
{
"model_name": "claude-3",
"litellm_params": {"model": "anthropic/claude-3-sonnet"},
}
]
mock_router.deployments_named.side_effect = {
"claude-3": (
{
"model_name": "claude-3",
"litellm_params": {"model": "anthropic/claude-3-sonnet"},
},
)
}.get
result = should_use_polling_for_request(
background_mode=True,
@ -1183,13 +1183,14 @@ class TestPollingConditionChecks:
)
mock_router = Mock()
mock_router.model_name_to_deployment_indices = {"deep-research": [0]}
mock_router.model_list = [
{
"model_name": "deep-research",
"litellm_params": {"model": "openai/o4-mini-deep-research"},
}
]
mock_router.deployments_named.side_effect = {
"deep-research": (
{
"model_name": "deep-research",
"litellm_params": {"model": "openai/o4-mini-deep-research"},
},
)
}.get
# Model alias "deep-research" is in native_background_mode
result = should_use_polling_for_request(

View file

@ -2,23 +2,35 @@ import json
from typing import Final
import pytest
from pydantic import ValidationError
from pydantic import BaseModel, ValidationError
from litellm.proxy._types import (
ROLES_WITHIN_ORG,
BudgetNewRequest,
ChangePasswordRequest,
GenerateKeyRequest,
GenerateKeyResponse,
KeyRequest,
LiteLLM_AuditLogs,
LiteLLM_BudgetTable,
LiteLLM_EndUserTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
LiteLLM_UserTable,
LitellmUserRoles,
NewCustomerRequest,
NewMCPServerRequest,
NewTeamRequest,
NewUserRequest,
NewUserResponse,
OrganizationMemberUpdateRequest,
RegenerateKeyRequest,
ResetSpendRequest,
UpdateKeyRequest,
UpdateMCPServerRequest,
UpdateTeamRequest,
UpdateUserRequest,
UpdateUserRequestNoUserIDorEmail,
UserAPIKeyAuth,
)
@ -213,6 +225,75 @@ def test_a_max_budget_that_was_actually_sent_is_kept(sent):
assert GenerateKeyRequest(max_budget=sent).max_budget == sent
RATE_LIMIT_REQUESTS: Final = (
(GenerateKeyRequest, {}),
(RegenerateKeyRequest, {}),
(UpdateKeyRequest, {"key": "sk-test"}),
(NewUserRequest, {}),
(UpdateUserRequestNoUserIDorEmail, {}),
(UpdateUserRequest, {"user_id": "user-test"}),
(NewTeamRequest, {}),
(UpdateTeamRequest, {"team_id": "team-test"}),
(BudgetNewRequest, {}),
(NewCustomerRequest, {"user_id": "customer-test"}),
)
@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS)
@pytest.mark.parametrize("field", ("rpm_limit", "tpm_limit"))
@pytest.mark.parametrize("value", (-1, -100))
def test_management_requests_reject_negative_rate_limits(
request_model: type[BaseModel], identifiers: dict[str, str], field: str, value: int
) -> None:
with pytest.raises(ValidationError) as error:
request_model.model_validate({**identifiers, field: value})
assert error.value.errors()[0]["loc"] == (field,)
@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS)
@pytest.mark.parametrize("value", (None, 0, 100))
def test_management_requests_preserve_explicit_rate_limits(
request_model: type[BaseModel], identifiers: dict[str, str], value: int | None
) -> None:
limits: Final = {"rpm_limit": value, "tpm_limit": value}
request: Final = request_model.model_validate({**identifiers, **limits})
assert request.model_dump(include={"rpm_limit", "tpm_limit"}, exclude_unset=True) == limits
@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS)
def test_management_requests_leave_omitted_rate_limits_unset(
request_model: type[BaseModel], identifiers: dict[str, str]
) -> None:
request: Final = request_model.model_validate({**identifiers})
assert request.model_dump(include={"rpm_limit", "tpm_limit"}) == {"rpm_limit": None, "tpm_limit": None}
assert request.model_dump(include={"rpm_limit", "tpm_limit"}, exclude_unset=True) == {}
@pytest.mark.parametrize(("response_model", "identifiers"), (
(GenerateKeyResponse, {"key": "sk-test"}),
(NewUserResponse, {"key": "sk-test"}),
(LiteLLM_TeamTable, {"team_id": "team-test"}),
(LiteLLM_UserTable, {"user_id": "user-test"}),
(LiteLLM_BudgetTable, {"budget_id": "budget-test"}),
))
def test_existing_negative_rate_limits_remain_readable(
response_model: type[BaseModel], identifiers: dict[str, str]
) -> None:
limits: Final = {"rpm_limit": -1, "tpm_limit": -100}
response: Final = response_model.model_validate({**identifiers, **limits})
assert response.model_dump(include={"rpm_limit", "tpm_limit"}) == limits
def test_existing_customer_negative_rate_limits_remain_readable() -> None:
customer: Final = LiteLLM_EndUserTable.model_validate({
"user_id": "customer-test", "blocked": False,
"litellm_budget_table": {"rpm_limit": -1, "tpm_limit": -100},
})
assert customer.litellm_budget_table is not None
assert customer.litellm_budget_table.rpm_limit == -1
assert customer.litellm_budget_table.tpm_limit == -100
USER_IDENTIFIER_CASES = (
({"user_id": "user-1"}, True),
({"user_email": "user@example.com"}, True),

View file

@ -1,4 +1,5 @@
import asyncio
import base64
import copy
import json
import os
@ -9208,6 +9209,97 @@ def test_default_team_settings_bool_turn_off_message_logging_redacts():
)
def _fake_chatgpt_oauth_token(account_id: str) -> str:
header = base64.urlsafe_b64encode(b'{"alg":"RS256"}').rstrip(b"=").decode()
payload = (
base64.urlsafe_b64encode(
json.dumps({"https://api.openai.com/auth": {"chatgpt_account_id": account_id}, "exp": 4102444800}).encode()
)
.rstrip(b"=")
.decode()
)
return f"{header}.{payload}.fake-signature"
CHATGPT_OAUTH_TOKEN = f"Bearer {_fake_chatgpt_oauth_token('acct-alice')}"
CHATGPT_ACCOUNT_ID_HEADER_NAME = "ChatGPT-Account-Id"
PLATFORM_API_KEY_BEARER = "Bearer sk-proj-fake-platform-key-0123456789"
CHATGPT_LEAK_TARGET_PROVIDERS = ["openai", "azure", "anthropic", "bedrock", "vertex_ai"]
def _codex_client_headers(
authorization_header_name: str = "authorization", account_id: str | None = "acct-alice"
) -> dict:
return {
"content-type": "application/json",
"originator": "codex_cli_rs",
authorization_header_name: CHATGPT_OAUTH_TOKEN,
**({CHATGPT_ACCOUNT_ID_HEADER_NAME: account_id} if account_id is not None else {}),
}
def test_clean_headers_keeps_chatgpt_oauth_bearer_when_gateway_key_is_separate():
result = clean_headers(
headers=Headers({**_codex_client_headers(), "x-litellm-api-key": "sk-proxy-virtual-key"}),
litellm_key_header_name="x-litellm-api-key",
forward_llm_provider_auth_headers=False,
authenticated_with_header="x-litellm-api-key",
)
assert result["authorization"] == CHATGPT_OAUTH_TOKEN
assert "x-litellm-api-key" not in result
def test_clean_headers_drops_bearer_that_authenticated_to_the_gateway():
result = clean_headers(
headers=Headers({"authorization": CHATGPT_OAUTH_TOKEN, "content-type": "application/json"}),
litellm_key_header_name=None,
forward_llm_provider_auth_headers=False,
authenticated_with_header="authorization",
)
assert _authorization_values(result) == []
@pytest.mark.parametrize("bearer", [PLATFORM_API_KEY_BEARER, "Bearer not.a-jwt", "Bearer a.bm90LWpzb24.c"])
def test_clean_headers_drops_non_chatgpt_bearers(bearer):
result = clean_headers(
headers=Headers({"authorization": bearer, "x-litellm-api-key": "sk-proxy-virtual-key"}),
litellm_key_header_name="x-litellm-api-key",
forward_llm_provider_auth_headers=False,
authenticated_with_header="x-litellm-api-key",
)
assert _authorization_values(result) == []
@pytest.mark.parametrize("authorization_header_name", AUTHORIZATION_HEADER_CASINGS)
def test_chatgpt_oauth_credential_and_account_reach_chatgpt_provider(authorization_header_name):
forwarded = _headers_forwarded_to(_codex_client_headers(authorization_header_name), "chatgpt")
assert _authorization_values(forwarded) == [CHATGPT_OAUTH_TOKEN]
assert forwarded[CHATGPT_ACCOUNT_ID_HEADER_NAME] == "acct-alice"
@pytest.mark.parametrize("custom_llm_provider", CHATGPT_LEAK_TARGET_PROVIDERS)
def test_chatgpt_oauth_credential_is_never_forwarded_to_other_providers(custom_llm_provider):
forwarded = _headers_forwarded_to(_codex_client_headers(), custom_llm_provider)
assert _authorization_values(forwarded) == []
assert CHATGPT_OAUTH_TOKEN not in forwarded.values()
assert CHATGPT_ACCOUNT_ID_HEADER_NAME not in forwarded
def test_chatgpt_account_header_is_not_forwarded_without_chatgpt_oauth_bearer():
headers = {**_codex_client_headers(), "authorization": PLATFORM_API_KEY_BEARER}
assert _headers_forwarded_to(headers, "chatgpt") == {}
def test_anthropic_oauth_bearer_does_not_reach_chatgpt_provider():
assert _authorization_values(_headers_forwarded_to(_client_headers(), "chatgpt")) == []
def test_add_user_api_key_auth_to_request_metadata_attributes_a_cli_session_to_its_alias():
data = {"model": "gpt-5.4-nano", "messages": [{"role": "user", "content": "hi"}], "litellm_metadata": {}}
session = UserAPIKeyAuth(

View file

@ -15207,21 +15207,23 @@ async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_coun
assert_loop_stayed_free(took, lags)
async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch):
from litellm.rust_bridge._native import Tokenizer
async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch: pytest.MonkeyPatch) -> None:
from queue import SimpleQueue
from threading import Thread, current_thread
from tokenizers import Tokenizer
from litellm import Router
from tests.unit.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags
claude_tokenizer: Final = litellm.utils.select_tokenizer("claude-fable-5")["tokenizer"]
claude_tokenizer: Final = Tokenizer.from_str(litellm.utils.claude_json_str)
loop_thread: Final = current_thread()
loader_threads: Final[SimpleQueue[Thread]] = SimpleQueue()
class SlowHubTokenizer:
@staticmethod
def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer:
time.sleep(0.3)
return claude_tokenizer
def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer:
loader_threads.put(current_thread())
return claude_tokenizer
monkeypatch.setattr("litellm.rust_bridge.tokenizer.from_pretrained", SlowHubTokenizer.from_pretrained)
monkeypatch.setattr("tokenizers.Tokenizer", types.SimpleNamespace(from_pretrained=from_pretrained))
monkeypatch.setattr(
"litellm.proxy.proxy_server.llm_router",
Router(
@ -15237,15 +15239,16 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp
),
)
response, took, lags = await timed_with_loop_lags(
lambda: proxy_server_module.token_counter(
litellm.utils._select_custom_tokenizer_helper.cache_clear()
try:
response: Final = await proxy_server_module.token_counter(
TokenCountRequest(model="self-hosted", prompt="count me off the loop")
)
)
assert response.tokenizer_type == "huggingface_tokenizer"
assert response.total_tokens > 0
assert_loop_stayed_free(took, lags)
assert loader_threads.get_nowait() is not loop_thread
assert response.tokenizer_type == "huggingface_tokenizer"
assert response.total_tokens > 0
finally:
litellm.utils._select_custom_tokenizer_helper.cache_clear()
async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revision_and_token(monkeypatch):

View file

@ -933,6 +933,42 @@ class TestResponsesInputToChatMessages:
) == [user_message]
class TestMergeProviderSpecificHeaders:
def test_scoped_headers_merge_only_for_matching_provider(self):
scoped = {"custom_llm_provider": "chatgpt", "extra_headers": {"Authorization": "Bearer client-token"}}
merged = ResponsesAPIRequestUtils.merge_provider_specific_headers(
extra_headers={"originator": "codex_cli_rs"},
provider_specific_header=scoped,
custom_llm_provider="chatgpt",
)
untouched = ResponsesAPIRequestUtils.merge_provider_specific_headers(
extra_headers={"originator": "codex_cli_rs"},
provider_specific_header=scoped,
custom_llm_provider="openai",
)
assert merged == {"originator": "codex_cli_rs", "Authorization": "Bearer client-token"}
assert untouched == {"originator": "codex_cli_rs"}
def test_sequence_of_scoped_entries_each_keep_their_own_scope(self):
entries = [
{"custom_llm_provider": "anthropic", "extra_headers": {"Authorization": "Bearer anthropic-token"}},
{"custom_llm_provider": "chatgpt", "extra_headers": {"Authorization": "Bearer chatgpt-token"}},
]
merged = ResponsesAPIRequestUtils.merge_provider_specific_headers(
extra_headers=None, provider_specific_header=entries, custom_llm_provider="chatgpt"
)
assert merged == {"Authorization": "Bearer chatgpt-token"}
def test_no_scoped_headers_returns_extra_headers_unchanged(self):
assert ResponsesAPIRequestUtils.merge_provider_specific_headers(
extra_headers=None, provider_specific_header=None, custom_llm_provider="chatgpt"
) is None
def test_responses_object_bills_an_explicit_cache_hit_at_the_explicit_rate():
model = "dashscope/qwen-responses-explicit-cache-test"
previous_row = litellm.model_cost.get(model)

View file

@ -112,6 +112,13 @@ def _heuristic_v2_artifact() -> TrainedTierArtifact:
)
def _answer_deployments_named(router: MagicMock) -> None:
"""Make a mock router answer `deployments_named` from its own `model_list`, the way Router answers from its index."""
router.deployments_named.side_effect = lambda name: (
tuple(deployment for deployment in router.model_list if deployment["model_name"] == name) or None
)
@pytest.fixture
def mock_router_instance():
"""Create a mock LiteLLM Router instance."""
@ -3181,7 +3188,7 @@ class TestCapabilityClassifier:
{"model_name": "efficient-model", "litellm_params": {"input_cost_per_token": 0.000001}},
{"model_name": "capable-model", "litellm_params": {"input_cost_per_token": 0.00001}},
]
mock_router_instance.model_name_to_deployment_indices = {"efficient-model": [0], "capable-model": [1]}
_answer_deployments_named(mock_router_instance)
router: Final = self._router(
mock_router_instance,
adaptive=True,
@ -5670,7 +5677,7 @@ class TestAdaptiveSoftFloors:
"model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}},
},
]
router.model_name_to_deployment_indices = {"cheap": [0], "premium": [1]}
_answer_deployments_named(router)
return router
@pytest.fixture
@ -6829,7 +6836,7 @@ class TestTierModelAffinity:
) -> ComplexityRouter:
mock_router_instance.cache = DualCache()
mock_router_instance.model_list = []
mock_router_instance.model_name_to_deployment_indices = {}
_answer_deployments_named(mock_router_instance)
return ComplexityRouter(
model_name="affinity-router",
litellm_router_instance=mock_router_instance,
@ -7547,7 +7554,7 @@ class TestSessionAffinity:
"model_info": {},
},
]
mock_router_instance.model_name_to_deployment_indices = {"cheap": [0]}
_answer_deployments_named(mock_router_instance)
router = ComplexityRouter(
model_name="hybrid",
litellm_router_instance=mock_router_instance,
@ -11351,7 +11358,7 @@ class TestClientHousekeepingCalls:
"model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}},
},
]
adaptive_instance.model_name_to_deployment_indices = {"cheap": [0], "premium": [1]}
_answer_deployments_named(adaptive_instance)
router = ComplexityRouter(
model_name="hybrid",
litellm_router_instance=adaptive_instance,
@ -13442,7 +13449,7 @@ class TestPlanModeTierFloor:
"model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}},
},
]
adaptive_instance.model_name_to_deployment_indices = {"cheap": [0], "premium": [1]}
_answer_deployments_named(adaptive_instance)
router = ComplexityRouter(
model_name="hybrid",
litellm_router_instance=adaptive_instance,
@ -14941,7 +14948,7 @@ class TestModalityRouting:
elif path == "adaptive_pick_rewritten":
config["adaptive"] = True
mock_router_instance.model_list = []
mock_router_instance.model_name_to_deployment_indices = {}
_answer_deployments_named(mock_router_instance)
router = self._router(mock_router_instance, config, vision)
result = await router.async_pre_routing_hook(model="m", request_kwargs=request_kwargs, messages=messages)
assert result.model == expected_model

View file

@ -107,8 +107,7 @@ class TestResolveBaseline:
served; a dashboard counterfactual must not be able to take routing down."""
class Exploding:
@property
def model_name_to_deployment_indices(self):
def deployments_named(self, model_name):
raise RuntimeError("router is mid-reload")
assert resolve_baseline(Exploding(), ["anything"]) is None

View file

@ -1,222 +0,0 @@
"""Unit tests for the Friendli transform in
`.github/scripts/auto_update_price_and_context_window_file.py`."""
from __future__ import annotations
import importlib.util
import sys
from pathlib import Path
import pytest
SCRIPT_PATH = (
Path(__file__).resolve().parents[2]
/ ".github"
/ "scripts"
/ "auto_update_price_and_context_window_file.py"
)
@pytest.fixture(scope="module")
def sync_module():
spec = importlib.util.spec_from_file_location(
"auto_update_price_and_context_window_file", SCRIPT_PATH
)
assert spec and spec.loader, f"Could not load spec for {SCRIPT_PATH}"
module = importlib.util.module_from_spec(spec)
sys.modules["auto_update_price_and_context_window_file"] = module
spec.loader.exec_module(module)
return module
def _reasoning_model(**overrides: object) -> dict:
model = {
"id": "zai-org/GLM-Test",
"base_model": "zhipuai/glm-test",
"context_length": 1048576,
"max_completion_tokens": 131072,
"pricing": {"input": "0.00000015", "output": "0.0000005", "input_cache_read": "0.00000003"},
"reasoning": True,
"reasoning_options": [{"type": "effort", "values": ["max", "high", "low"]}],
"functionality": {
"tool_call": True,
"parallel_tool_call": True,
"structured_output": True,
"system_messages": True,
"tool_choice": True,
},
"input_modalities": ["text", "image", "video"],
"mode": "chat",
}
model.update(overrides)
return model
def test_transform_emits_declared_effort_levels_in_canonical_order(sync_module):
entry = sync_module.transform_friendli_data([_reasoning_model()], {})[
"friendliai/zai-org/GLM-Test"
]
assert entry["supports_reasoning"] is True
assert entry["reasoning_effort_levels"] == ["low", "high", "max"]
assert not any(k.endswith("_reasoning_effort") for k in entry)
def test_transform_reasoning_model_without_effort_options_declares_empty_levels(sync_module):
model = _reasoning_model(reasoning_options=[{"type": "budget_tokens", "values": []}])
entry = sync_module.transform_friendli_data([model], {})["friendliai/zai-org/GLM-Test"]
assert entry["reasoning_effort_levels"] == []
def test_transform_non_reasoning_model_declares_no_levels(sync_module):
model = _reasoning_model(reasoning=False, reasoning_options=[])
entry = sync_module.transform_friendli_data([model], {})["friendliai/zai-org/GLM-Test"]
assert entry["supports_reasoning"] is False
assert "reasoning_effort_levels" not in entry
def test_transform_max_tokens_mirrors_output_cap_not_context(sync_module):
entry = sync_module.transform_friendli_data([_reasoning_model()], {})[
"friendliai/zai-org/GLM-Test"
]
assert entry["max_input_tokens"] == 1048576
assert entry["max_output_tokens"] == 131072
assert entry["max_tokens"] == entry["max_output_tokens"]
def test_transform_prompt_caching_follows_cache_pricing(sync_module):
cached = sync_module.transform_friendli_data([_reasoning_model()], {})[
"friendliai/zai-org/GLM-Test"
]
assert cached["supports_prompt_caching"] is True
assert cached["cache_read_input_token_cost"] == 3e-08
uncached_model = _reasoning_model(pricing={"input": "0.00000014", "output": "0.0000004"})
uncached = sync_module.transform_friendli_data([uncached_model], {})[
"friendliai/zai-org/GLM-Test"
]
assert uncached["supports_prompt_caching"] is False
assert "cache_read_input_token_cost" not in uncached
def test_transform_modalities_set_vision_image_and_video_flags(sync_module):
entry = sync_module.transform_friendli_data([_reasoning_model()], {})[
"friendliai/zai-org/GLM-Test"
]
assert entry["supports_vision"] is True
assert entry["supports_image_input"] is True
assert entry["supports_video_input"] is True
text_only = _reasoning_model(input_modalities=["text"])
entry_text = sync_module.transform_friendli_data([text_only], {})[
"friendliai/zai-org/GLM-Test"
]
assert entry_text["supports_vision"] is False
assert entry_text["supports_image_input"] is False
assert entry_text["supports_video_input"] is False
def test_transform_skips_rows_without_valid_token_prices_so_priced_local_entries_survive(sync_module):
local = {
"friendliai/zai-org/GLM-Test": {
"litellm_provider": "friendliai",
"input_cost_per_token": 1.5e-07,
"output_cost_per_token": 5e-07,
}
}
unpriced_rows = [
_reasoning_model(pricing={}),
_reasoning_model(pricing=None),
_reasoning_model(pricing={"input": "0.00000015"}),
_reasoning_model(pricing={"output": "0.0000005"}),
_reasoning_model(pricing={"input": "not-a-number", "output": "0.0000005"}),
_reasoning_model(pricing={"input": "-0.00000015", "output": "0.0000005"}),
_reasoning_model(pricing={"input": "inf", "output": "0.0000005"}),
_reasoning_model(pricing={"input": "nan", "output": "0.0000005"}),
]
remote = sync_module.transform_friendli_data(unpriced_rows, local)
assert remote == {}
sync_module.sync_local_data_with_remote(local, remote, replace_keys=frozenset(remote))
assert local["friendliai/zai-org/GLM-Test"]["input_cost_per_token"] == 1.5e-07
assert local["friendliai/zai-org/GLM-Test"]["output_cost_per_token"] == 5e-07
def test_transform_keeps_zero_priced_rows(sync_module):
free_model = _reasoning_model(pricing={"input": "0", "output": "0"})
entry = sync_module.transform_friendli_data([free_model], {})["friendliai/zai-org/GLM-Test"]
assert entry["input_cost_per_token"] == 0.0
assert entry["output_cost_per_token"] == 0.0
def test_transforms_survive_failed_fetch(sync_module):
assert sync_module.transform_friendli_data(None, {}) == {}
assert sync_module.transform_friendli_data([], {}) == {}
assert sync_module.transform_openrouter_data(None) == {}
assert sync_module.transform_vercel_ai_gateway_data(None) == {}
def test_vercel_transform_skips_rows_without_token_pricing_or_limits(sync_module):
rows = [
{
"id": "wan-video",
"pricing": {"video_duration_pricing": [{"resolution": "720p", "cost_per_second": "0.1"}]},
},
{
"id": "qwen3-embedding",
"context_window": 32768,
"max_tokens": 32768,
"pricing": {"input": "0.00000001"},
},
{
"id": "no-limits-chat",
"pricing": {"input": "0.000001", "output": "0.000002"},
},
{
"id": "good-chat",
"context_window": 128000,
"max_tokens": 8192,
"pricing": {"input": "0.000001", "output": "0.000002"},
},
]
transformed = sync_module.transform_vercel_ai_gateway_data(rows)
assert list(transformed) == ["vercel_ai_gateway/good-chat"]
assert transformed["vercel_ai_gateway/good-chat"]["input_cost_per_token"] == 1e-06
assert transformed["vercel_ai_gateway/good-chat"]["output_cost_per_token"] == 2e-06
def test_sync_replaces_friendli_entries_so_dropped_cache_pricing_does_not_survive(sync_module):
local = {
"friendliai/zai-org/GLM-Test": {
"litellm_provider": "friendliai",
"cache_read_input_token_cost": 3e-08,
"supports_prompt_caching": True,
}
}
uncached_model = _reasoning_model(pricing={"input": "0.00000014", "output": "0.0000004"})
remote = sync_module.transform_friendli_data([uncached_model], local)
sync_module.sync_local_data_with_remote(local, remote, replace_keys=frozenset(remote))
synced = local["friendliai/zai-org/GLM-Test"]
assert "cache_read_input_token_cost" not in synced
assert synced["supports_prompt_caching"] is False
def test_sync_still_merges_entries_outside_replace_keys(sync_module):
local = {"openrouter/some-model": {"input_cost_per_token": 1e-06, "supports_vision": True}}
remote = {"openrouter/some-model": {"input_cost_per_token": 2e-06}}
sync_module.sync_local_data_with_remote(local, remote)
assert local["openrouter/some-model"] == {"input_cost_per_token": 2e-06, "supports_vision": True}
def test_transform_inherits_allowlisted_keys_from_base_model_entry(sync_module):
local = {
"zhipuai/glm-test": {
"supports_pdf_input": True,
"supports_assistant_prefill": True,
"input_cost_per_token": 9e-06,
}
}
entry = sync_module.transform_friendli_data([_reasoning_model()], local)[
"friendliai/zai-org/GLM-Test"
]
assert entry["supports_pdf_input"] is True
assert entry["supports_assistant_prefill"] is True
assert entry["input_cost_per_token"] == 1.5e-07

View file

@ -27,6 +27,23 @@ def prices() -> dict:
return json.loads(PRICES_PATH.read_text())
@pytest.mark.parametrize("model", ("azure/gpt-audio-1.5-2026-02-23", "azure/gpt-audio-mini-2025-10-06"))
def test_azure_audio_snapshot_capabilities_match_declared_modalities(model: str, local_model_cost_map: None) -> None:
catalog: Final = TypeAdapter(dict[str, JsonValue]).validate_json(PRICES_PATH.read_bytes())
backup: Final = TypeAdapter(dict[str, JsonValue]).validate_json(BACKUP_PRICES_PATH.read_bytes())
entry: Final = catalog[model]
assert isinstance(entry, dict)
assert entry == backup[model]
input_modalities: Final = entry["supported_modalities"]
output_modalities: Final = entry["supported_output_modalities"]
assert isinstance(input_modalities, list)
assert isinstance(output_modalities, list)
expected: Final = ("audio" in input_modalities, "audio" in output_modalities)
info: Final = litellm.get_model_info(model)
assert (info["supports_audio_input"], info["supports_audio_output"]) == expected
assert (litellm.supports_audio_input(model), litellm.supports_audio_output(model)) == expected
class _CopilotEndpointRow(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")

View file

@ -5,6 +5,7 @@ import functools
import gc
import itertools
import json
import linecache
import logging
import os
import sys
@ -13,8 +14,8 @@ import warnings
from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping
from datetime import datetime, timedelta, timezone
from pathlib import Path
from types import MappingProxyType, SimpleNamespace
from typing import TYPE_CHECKING, Final, Literal, cast
from types import CodeType, FrameType, MappingProxyType, SimpleNamespace
from typing import TYPE_CHECKING, Final, Literal, TypeVar, cast
from unittest.mock import AsyncMock, MagicMock, Mock, patch
import httpx
@ -60,7 +61,10 @@ from litellm.router import (
_without_line_breaks,
Span,
)
from litellm.proxy.response_polling.polling_handler import should_use_polling_for_request
from litellm.router_strategy import simple_shuffle
from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter
from litellm.router_strategy.savings_baseline import Baseline, _models_in
from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit
from litellm.router_utils.cooldown_handlers import (
_async_get_cooldown_deployments,
@ -279,6 +283,15 @@ def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata():
assert kwargs["litellm_metadata"] == {"baz": 123}
def test_router_model_group_llm_provider_resolves_chatgpt_without_login(tmp_path, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
params: Final = LiteLLM_Params(model="chatgpt/gpt-5.4")
assert Router._model_group_llm_provider(params) == ("gpt-5.4", "chatgpt")
assert Router._model_group_llm_provider(LiteLLM_Params(model="no-such-provider-model-xyz")) == ("", "")
assert not any(tmp_path.iterdir())
def test_router_with_model_info_and_model_group():
"""
Test edge case where user specifies model_group in model_info
@ -30004,6 +30017,319 @@ async def test_responses_fallback_wrapper_disconnect_billing_keeps_the_primary_s
assert wrapped.partial_response_on_disconnect().response.usage.output_tokens == primary_tokens + fallback_tokens
_SYNC_SECONDS: Final = 10
_ROUTER_FILE: Final = Router.__init__.__code__.co_filename
_INDEX_LOCK_LINES: Final = frozenset(
number
for number, line in enumerate(linecache.getlines(_ROUTER_FILE), start=1)
if "with self._deployment_index_lock" in line
)
_Read = TypeVar("_Read")
def _concurrent_delete_deployment(index: int) -> dict[str, object]:
return {
"model_name": f"model-{index}",
"litellm_params": {"model": f"openai/gpt-{index}", "api_key": "k"},
"model_info": {
"id": f"deployment-{index}",
"team_id": "team-1",
"team_public_model_name": "public-shared" if index >= 2 else f"public-{index}",
},
}
def _concurrent_delete_router() -> Router:
return Router(model_list=[_concurrent_delete_deployment(index) for index in range(4)])
def _read_during_held_write(
read: Callable[[], _Read],
write: Callable[[], object],
event: Literal["call", "c_call"] = "call",
write_after: CodeType | None = None,
) -> _Read:
"""Run `read` on its own thread while `write` waits at its first `event` in the re-index, after the list shrank and
before the maps caught up. With `write_after`, the write starts when `read` returns from that function. The write
resumes once `read` returns or reaches the router's index lock, so a read that skips the lock meets the half-updated
list and a read that takes the lock waits for the write"""
reindex: Final = Router._update_deployment_indices_after_removal.__code__
held: Final = threading.Event()
release: Final = threading.Event()
settled: Final = threading.Event()
released_in_time: Final[list[bool]] = []
outcome: Final[list[_Read]] = []
failure: Final[list[BaseException]] = []
def hold_in_reindex(frame: FrameType, observed: str, _arg: object) -> None:
if observed == event and frame.f_code is reindex and not held.is_set():
held.set()
released_in_time.append(release.wait(timeout=_SYNC_SECONDS))
def run_write() -> None:
previous: Final = sys.getprofile()
sys.setprofile(hold_in_reindex)
try:
write()
finally:
sys.setprofile(previous)
writer: Final = threading.Thread(target=run_write)
def start_write() -> None:
writer.start()
assert held.wait(timeout=_SYNC_SECONDS), "the write never reached re-indexing"
def watch_read(frame: FrameType, observed: str, _arg: object) -> Callable[..., object] | None:
if frame.f_code.co_filename != _ROUTER_FILE:
return None
if observed == "return" and frame.f_code is write_after and writer.ident is None:
start_write()
elif observed == "line" and held.is_set() and frame.f_lineno in _INDEX_LOCK_LINES:
settled.set()
return watch_read
def run_read() -> None:
previous: Final = sys.gettrace()
sys.settrace(watch_read)
try:
outcome.append(read())
except BaseException as error:
failure.append(error)
finally:
sys.settrace(previous)
settled.set()
reader: Final = threading.Thread(target=run_read)
try:
if write_after is None:
start_write()
reader.start()
assert settled.wait(timeout=_SYNC_SECONDS), "the read neither returned nor reached the index lock"
finally:
release.set()
if reader.ident is not None:
reader.join()
if writer.ident is not None:
writer.join()
assert released_in_time == [True], "the write was not held for the whole read"
if failure:
raise failure[0]
return outcome[0]
def _deployment_ids(deployments: list[DeploymentTypedDict]) -> tuple[str, ...]:
return tuple(deployment["model_info"]["id"] for deployment in deployments)
@pytest.mark.parametrize(
("read", "expected"),
[
pytest.param(
lambda router: router.get_deployment("deployment-3").model_info.id, "deployment-3", id="get_deployment"
),
pytest.param(
lambda router: router.get_model_info("deployment-3")["model_info"]["id"],
"deployment-3",
id="get_model_info",
),
pytest.param(
lambda router: _deployment_ids(router.get_all_deployments("model-3")),
("deployment-3",),
id="get_all_deployments",
),
pytest.param(
lambda router: _deployment_ids(router.get_all_deployments("public-shared", team_id="team-1")),
("deployment-2", "deployment-3"),
id="get_all_deployments-team",
),
pytest.param(
lambda router: router.map_team_model("public-shared", "team-1"), "public-shared", id="map_team_model"
),
pytest.param(lambda router: router.get_model_ids("model-3"), ["deployment-3"], id="get_model_ids-by-name"),
pytest.param(
lambda router: router.get_model_ids(), ["deployment-1", "deployment-2", "deployment-3"], id="get_model_ids"
),
pytest.param(
lambda router: router.get_deployment_by_model_group_name("model-3").model_info.id,
"deployment-3",
id="get_deployment_by_model_group_name",
),
pytest.param(
lambda router: router.get_model_listing_info("model-3").cost_map_keys,
("openai/gpt-3",),
id="get_model_listing_info",
),
pytest.param(
lambda router: router.team_model_has_alternatives("deployment-3"), True, id="team_model_has_alternatives"
),
pytest.param(
lambda router: (
router._get_model_group_deployment_usable_by_team(
model_group_name="model-3", team_id="team-1"
).model_info.id
),
"deployment-3",
id="model_group_deployment_usable_by_team",
),
pytest.param(
lambda router: (
router._get_team_public_name_deployment(model_id="public-shared", team_id="team-1").model_info.id
),
"deployment-2",
id="team_public_name_deployment",
),
pytest.param(
lambda router: _deployment_ids(router._team_deployments_across_teams("public-shared")),
("deployment-2", "deployment-3"),
id="team_deployments_across_teams",
),
pytest.param(
lambda router: _deployment_ids(list(router.deployments_named("model-3") or ())),
("deployment-3",),
id="deployments_named",
),
pytest.param(
lambda router: should_use_polling_for_request(
background_mode=True,
polling_via_cache_enabled=["openai"],
redis_cache=MagicMock(),
model="model-3",
llm_router=router,
),
True,
id="response_polling_provider_lookup",
),
pytest.param(
lambda router: _models_in(router, "model-3"),
(Baseline("openai/gpt-3", "deployment-3"),),
id="savings_baseline_candidates",
),
],
)
def test_lookup_during_a_concurrent_delete_returns_the_requested_deployment(
read: Callable[[Router], object], expected: object
) -> None:
router: Final = _concurrent_delete_router()
observed: Final = _read_during_held_write(lambda: read(router), lambda: router.delete_deployment("deployment-0"))
assert observed == expected
def test_lookup_during_an_upsert_returns_the_requested_deployment() -> None:
router: Final = _concurrent_delete_router()
rotated: Final = Deployment(
**{**_concurrent_delete_deployment(0), "litellm_params": {"model": "openai/gpt-0", "api_key": "rotated"}}
)
observed: Final = _read_during_held_write(
lambda: router.get_deployment("deployment-3").model_info.id, lambda: router.upsert_deployment(rotated)
)
assert observed == "deployment-3"
def test_add_during_a_concurrent_delete_indexes_the_added_deployment() -> None:
router: Final = _concurrent_delete_router()
_read_during_held_write(
lambda: router.add_deployment(Deployment(**_concurrent_delete_deployment(4))),
lambda: router.delete_deployment("deployment-0"),
)
assert [router.get_deployment(f"deployment-{index}").model_info.id for index in range(1, 5)] == [
f"deployment-{index}" for index in range(1, 5)
]
def test_set_model_list_during_a_concurrent_delete_indexes_the_new_list() -> None:
router: Final = _concurrent_delete_router()
_read_during_held_write(
lambda: router.set_model_list([_concurrent_delete_deployment(index) for index in range(10, 14)]),
lambda: router.delete_deployment("deployment-0"),
event="c_call",
)
assert router.get_model_ids() == [f"deployment-{index}" for index in range(10, 14)]
def test_team_model_has_alternatives_when_a_delete_starts_after_its_deployment_lookup() -> None:
router: Final = _concurrent_delete_router()
has_alternatives: Final = _read_during_held_write(
lambda: router.team_model_has_alternatives("deployment-3"),
lambda: router.delete_deployment("deployment-0"),
write_after=Router.get_deployment.__code__,
)
assert has_alternatives is True
def test_init_adaptive_router_deployment_during_a_concurrent_delete_prices_the_listed_model() -> None:
priced: Final = {
**_concurrent_delete_deployment(3),
"litellm_params": {"model": "openai/gpt-3", "api_key": "k", "input_cost_per_token": 3e-06},
}
router: Final = Router(model_list=[*(_concurrent_delete_deployment(index) for index in range(3)), priced])
adaptive: Final = Deployment(
model_name="adaptive",
litellm_params=LiteLLM_Params(
model="auto_router/adaptive_router", adaptive_router_config={"available_models": ["model-3"]}
),
model_info=ModelInfo(id="adaptive"),
)
_read_during_held_write(
lambda: router.init_adaptive_router_deployment(deployment=adaptive),
lambda: router.delete_deployment("deployment-0"),
)
assert router.adaptive_routers["adaptive"][0].strategy.model_to_cost == {"model-3": 3e-06}
def test_complexity_router_adaptive_pool_during_a_concurrent_delete_prices_the_listed_model() -> None:
priced: Final = {
**_concurrent_delete_deployment(3),
"litellm_params": {"model": "openai/gpt-3", "api_key": "k", "input_cost_per_token": 3e-06},
}
router: Final = Router(model_list=[*(_concurrent_delete_deployment(index) for index in range(3)), priced])
complexity: Final = ComplexityRouter(
model_name="hybrid",
litellm_router_instance=router,
complexity_router_config={"adaptive": True, "tiers": {"SIMPLE": ["model-3"]}},
)
adaptive: Final = _read_during_held_write(
complexity._ensure_adaptive_router, lambda: router.delete_deployment("deployment-0")
)
assert adaptive is not None
assert adaptive.model_to_cost == {"model-3": 3e-06}
def test_delete_of_an_unknown_deployment_returns_none_and_keeps_the_list() -> None:
router: Final = _concurrent_delete_router()
assert router.delete_deployment("deployment-9") is None
assert router.get_model_ids() == [f"deployment-{index}" for index in range(4)]
def test_get_model_ids_by_name_skips_team_deployments_only_when_asked() -> None:
shared: Final = {"model_name": "shared", "litellm_params": {"model": "openai/gpt-shared", "api_key": "k"}}
router: Final = Router(
model_list=[
{**shared, "model_info": {"id": "global"}},
{**shared, "model_info": {"id": "team", "team_id": "team-1", "team_public_model_name": "public"}},
]
)
assert router.get_model_ids("shared", exclude_team_models=True) == ["global"]
assert router.get_model_ids("shared") == ["global", "team"]
_IMAGE_PART = {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}
_IMAGE_CHAT = [{"role": "user", "content": [{"type": "text", "text": "What is in this image?"}, _IMAGE_PART]}]
_IMAGE_RESPONSES_INPUT = [
@ -30153,7 +30479,7 @@ def test_request_carries_image_sees_chat_parts_and_responses_input_items():
@pytest.mark.parametrize(
("model", "expected_provider"),
[("cloudflare/@cf/cloudflare/clef-flash", "cloudflare"), ("dotprompt/clef", None)],
[("cloudflare/@cf/cloudflare/clef-flash", "cloudflare"), ("dotprompt/clef", None), ("chatgpt/gpt-5.4", "chatgpt")],
)
def test_resolve_deployment_provider_accepts_what_the_router_loads(model: str, expected_provider: str | None):
resolved: Final = Router.resolve_deployment_provider(LiteLLM_Params(model=model))

View file

@ -1,378 +0,0 @@
import importlib.util
import json
from pathlib import Path
from types import MappingProxyType
import pytest
ROOT = Path(__file__).resolve().parents[2]
SCRIPT = ROOT / "scripts" / "sync_together_ai_models.py"
FIXTURES = Path(__file__).resolve().parent / "fixtures" / "together_ai_sync"
_spec = importlib.util.spec_from_file_location("sync_together_ai_models", SCRIPT)
assert _spec is not None and _spec.loader is not None
sync = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(sync)
RECORDED_CATALOG = sync.load_catalog(FIXTURES.joinpath("models_serverless.json").read_bytes())
RECORDED_DOC = sync.parse_deprecations(FIXTURES.joinpath("deprecations.md").read_text())
def _doc(removal_dates: dict[str, str], redirects: dict[str, str] | None = None) -> object:
return sync.DeprecationDoc(
removal_dates=MappingProxyType(removal_dates),
redirects=MappingProxyType(redirects or {}),
)
def _chat_model(model_id: str, ctx: int = 4096, price: float = 1.0, cached: float | None = None) -> object:
return sync.CatalogModel(
id=model_id,
type="chat",
context_length=ctx,
pricing=sync.CatalogPricing(input=price, output=price, cached_input=cached),
)
@pytest.mark.parametrize(
("per_million", "expected"),
[
(3, 3e-06),
(15, 1.5e-05),
(1.4, 1.4e-06),
(0.25999999999999995, 2.6e-07),
(0.060000000000000005, 6e-08),
(1.0399999999999998, 1.04e-06),
(0, 0.0),
],
)
def test_per_token_normalizes_float_artifacts(per_million: float, expected: float) -> None:
assert sync.per_token(per_million) == expected
def test_parse_deprecations_recorded_fixture() -> None:
assert dict(RECORDED_DOC.redirects) == {
"mistralai/Mistral-7B-Instruct-v0.3": "mistralai/Ministral-3-14B-Instruct-2512",
"Kimi-K2": "Kimi-K2-0905",
"DeepSeek-V3": "DeepSeek-V3.1",
"DeepSeek-V3-0324": "DeepSeek-V3.1",
"DeepSeek-R1": "DeepSeek-R1-0528",
}
assert len(RECORDED_DOC.removal_dates) == 208
assert RECORDED_DOC.removal_dates["google/gemma-3n-E4B-it"] == "2026-08-04"
def test_parse_deprecations_duplicate_rows_keep_most_recent_date() -> None:
assert RECORDED_DOC.removal_dates["Qwen/Qwen3-235B-A22B-Thinking-2507"] == "2026-04-16"
@pytest.mark.parametrize(
"markdown",
[
"# Deprecations\n\nNothing here anymore.\n",
"\n## Active model redirects\n\n| A | B |\n| --- | --- |\n| `x` | `y` |\n\n## Something else\n",
"\n## Deprecation history\n\n### Inference\n\n| Date | Model | R |\n| --- | --- | --- |\n| 2026-01-01 | `m` | No |\n",
],
)
def test_parse_deprecations_raises_when_a_table_parses_empty(markdown: str) -> None:
with pytest.raises(sync.SyncError):
sync.parse_deprecations(markdown)
def test_load_catalog_raises_on_shape_change() -> None:
with pytest.raises(sync.SyncError):
sync.load_catalog(b'[{"id": "x", "type": "chat"}]')
def test_load_catalog_raises_when_no_token_models_remain() -> None:
only_video = json.dumps([{"id": "v", "type": "video", "pricing": {"input": 0, "output": 0}}]).encode()
with pytest.raises(sync.SyncError):
sync.load_catalog(only_video)
def test_recorded_catalog_counts() -> None:
assert len(RECORDED_CATALOG) == 102
assert sum(1 for model in RECORDED_CATALOG if model.type in sync.TYPE_TO_MODE) == 26
assert sum(1 for model in RECORDED_CATALOG if model.pricing.cached_input) == 13
def test_added_chat_model_matches_reviewed_registry_shape() -> None:
outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC)
assert len(outcome.added) == 26
assert not outcome.deprecated
assert outcome.cost_map["together_ai/moonshotai/Kimi-K3"] == {
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 1048576,
"max_tokens": 1048576,
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"source": "https://docs.together.ai/docs/serverless-models",
"supports_function_calling": True,
"supports_parallel_function_calling": True,
"supports_prompt_caching": True,
"supports_reasoning": True,
"supports_response_schema": True,
"supports_tool_choice": True,
"supports_vision": True,
}
def test_added_embedding_model_has_no_output_token_cap() -> None:
outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC)
assert outcome.cost_map["together_ai/intfloat/multilingual-e5-large-instruct"] == {
"input_cost_per_token": 2e-08,
"litellm_provider": "together_ai",
"max_input_tokens": 514,
"max_tokens": 514,
"mode": "embedding",
"output_cost_per_token": 2e-08,
"output_vector_size": 1024,
"source": "https://docs.together.ai/docs/serverless-models",
}
def test_moderation_type_maps_to_chat_mode() -> None:
outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC)
guard = outcome.cost_map["together_ai/meta-llama/Llama-Guard-4-12B"]
assert guard["mode"] == "chat"
assert "max_output_tokens" not in guard
def test_output_ceiling_comes_from_the_rule_never_from_context_length() -> None:
glm = next(model for model in RECORDED_CATALOG if model.id == "zai-org/GLM-5.2")
fresh = sync.compute_sync({}, [_chat_model("acme/unreviewed", ctx=1048576), glm], _doc({"x": "2026-01-01"}))
unreviewed = fresh.cost_map["together_ai/acme/unreviewed"]
assert "max_output_tokens" not in unreviewed
assert (unreviewed["max_input_tokens"], unreviewed["max_tokens"]) == (1048576, 1048576)
reviewed = fresh.cost_map["together_ai/zai-org/GLM-5.2"]
assert (reviewed["max_input_tokens"], reviewed["max_output_tokens"], reviewed["max_tokens"]) == (
1048575,
128000,
128000,
)
inflated = {
"together_ai/zai-org/GLM-5.2": {
"input_cost_per_token": 1.4e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 1048575,
"max_output_tokens": 1048575,
"max_tokens": 1048575,
"mode": "chat",
"output_cost_per_token": 4.4e-06,
}
}
corrected = sync.compute_sync(inflated, [glm], _doc({"x": "2026-01-01"}))
assert corrected.cost_map["together_ai/zai-org/GLM-5.2"]["max_output_tokens"] == 128000
assert any("max_output_tokens: 1048575 -> 128000" in line for line in corrected.updated)
def test_docs_removed_but_live_model_stays_live_with_warning() -> None:
outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC)
gemma = outcome.cost_map["together_ai/google/gemma-3n-E4B-it"]
assert "deprecation_date" not in gemma
assert any("gemma-3n-E4B-it" in warning and "2026-08-04" in warning for warning in outcome.warnings)
def test_price_change_updates_api_fields_and_keeps_curated_ones() -> None:
registry = {
"together_ai/acme/chat-1": {
"input_cost_per_token": 9e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 4096,
"max_output_tokens": 2048,
"max_tokens": 2048,
"mode": "chat",
"output_cost_per_token": 9e-07,
"supports_audio_input": True,
}
}
outcome = sync.compute_sync(registry, [_chat_model("acme/chat-1", ctx=8192, price=2.0)], _doc({"x": "2026-01-01"}))
entry = outcome.cost_map["together_ai/acme/chat-1"]
assert entry["input_cost_per_token"] == 2e-06
assert entry["max_input_tokens"] == 8192
assert entry["max_output_tokens"] == 2048
assert entry["supports_audio_input"] is True
assert len(outcome.updated) == 1
assert "input_cost_per_token" in outcome.updated[0]
def test_cached_input_appearing_and_disappearing() -> None:
registry = {
"together_ai/acme/chat-1": {
"cache_read_input_token_cost": 1e-07,
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 1e-06,
"supports_prompt_caching": True,
},
"together_ai/acme/chat-2": {
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 1e-06,
},
}
catalog = [_chat_model("acme/chat-1"), _chat_model("acme/chat-2", cached=0.25999999999999995)]
outcome = sync.compute_sync(registry, catalog, _doc({"x": "2026-01-01"}))
assert "cache_read_input_token_cost" not in outcome.cost_map["together_ai/acme/chat-1"]
assert "supports_prompt_caching" not in outcome.cost_map["together_ai/acme/chat-1"]
assert outcome.cost_map["together_ai/acme/chat-2"]["cache_read_input_token_cost"] == 2.6e-07
assert outcome.cost_map["together_ai/acme/chat-2"]["supports_prompt_caching"] is True
def test_capability_rule_backfills_existing_entry() -> None:
registry = {
"together_ai/moonshotai/Kimi-K3": {
"input_cost_per_token": 3e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 1048576,
"mode": "chat",
"output_cost_per_token": 1.5e-05,
}
}
kimi = next(model for model in RECORDED_CATALOG if model.id == "moonshotai/Kimi-K3")
outcome = sync.compute_sync(registry, [kimi], _doc({"x": "2026-01-01"}))
assert outcome.cost_map["together_ai/moonshotai/Kimi-K3"]["supports_reasoning"] is True
assert any("supports_reasoning" in line for line in outcome.updated)
def test_disappeared_model_gets_docs_date_and_is_never_deleted() -> None:
registry = {
"together_ai/acme/gone": {
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",
"mode": "chat",
"output_cost_per_token": 1e-06,
}
}
outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"acme/gone": "2026-07-01"}))
assert outcome.cost_map["together_ai/acme/gone"]["deprecation_date"] == "2026-07-01"
assert outcome.deprecated == ("together_ai/acme/gone: deprecation_date",)
def test_disappeared_model_without_docs_date_warns_instead() -> None:
registry = {
"together_ai/acme/gone": {
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",
"mode": "chat",
"output_cost_per_token": 1e-06,
}
}
outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"other": "2026-07-01"}))
assert "deprecation_date" not in outcome.cost_map["together_ai/acme/gone"]
assert not outcome.deprecated
assert any("acme/gone" in warning and "human" in warning for warning in outcome.warnings)
def test_curated_deprecation_date_is_never_overwritten() -> None:
registry = {
"together_ai/acme/gone": {
"deprecation_date": "2026-06-15",
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",
"mode": "chat",
"output_cost_per_token": 1e-06,
}
}
outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"acme/gone": "2026-07-01"}))
assert outcome.cost_map["together_ai/acme/gone"]["deprecation_date"] == "2026-06-15"
assert any("2026-06-15" in warning and "2026-07-01" in warning for warning in outcome.warnings)
def test_redirect_chain_resolves_to_final_live_model() -> None:
doc = _doc({"acme/a": "2026-01-01"}, redirects={"acme/a": "acme/b", "acme/b": "acme/c"})
live = frozenset({"acme/c"})
assert sync.resolve_successor("acme/a", doc, live) == "acme/c"
def test_redirect_dead_end_yields_no_successor() -> None:
doc = _doc({"acme/a": "2026-01-01"}, redirects={"acme/a": "acme/b"})
assert sync.resolve_successor("acme/a", doc, frozenset({"acme/other"})) is None
def test_redirect_short_names_resolve_by_unique_suffix() -> None:
doc = _doc({"moonshotai/Kimi-K2": "2026-01-01"}, redirects={"Kimi-K2": "Kimi-K2-0905"})
live = frozenset({"moonshotai/Kimi-K2-0905"})
assert sync.resolve_successor("moonshotai/Kimi-K2", doc, live) == "moonshotai/Kimi-K2-0905"
def test_successor_written_only_when_not_curated() -> None:
registry = {
"together_ai/acme/a": {
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",
"mode": "chat",
"output_cost_per_token": 1e-06,
},
"together_ai/acme/b": {
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",
"metadata": {"successor": "together_ai/acme/curated"},
"mode": "chat",
"output_cost_per_token": 1e-06,
},
}
doc = _doc({"acme/a": "2026-01-01", "acme/b": "2026-01-01"}, redirects={"acme/a": "acme/c", "acme/b": "acme/c"})
outcome = sync.compute_sync(registry, [_chat_model("acme/c")], doc)
assert outcome.cost_map["together_ai/acme/a"]["metadata"] == {"successor": "together_ai/acme/c"}
assert outcome.cost_map["together_ai/acme/b"]["metadata"] == {"successor": "together_ai/acme/curated"}
assert any("acme/curated" in warning for warning in outcome.warnings)
def test_reappearance_clears_deprecation_date() -> None:
registry = {
"together_ai/acme/back": {
"deprecation_date": "2026-05-01",
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 1e-06,
}
}
outcome = sync.compute_sync(registry, [_chat_model("acme/back")], _doc({"x": "2026-01-01"}))
assert "deprecation_date" not in outcome.cost_map["together_ai/acme/back"]
assert outcome.reappeared == ("together_ai/acme/back",)
def test_new_chat_model_without_rule_is_flagged() -> None:
outcome = sync.compute_sync({}, [_chat_model("acme/unreviewed")], _doc({"x": "2026-01-01"}))
assert any("acme/unreviewed" in warning and "capability rule" in warning for warning in outcome.warnings)
def test_new_keys_land_at_the_end_of_the_provider_block() -> None:
registry = {
"aaa": {"mode": "chat"},
"together_ai/acme/old": {
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",
"mode": "chat",
"output_cost_per_token": 1e-06,
},
"zzz": {"mode": "chat"},
}
outcome = sync.compute_sync(registry, [_chat_model("acme/old"), _chat_model("acme/new")], _doc({"x": "2026-01-01"}))
assert list(outcome.cost_map) == ["aaa", "together_ai/acme/old", "together_ai/acme/new", "zzz"]
def test_sync_is_idempotent_over_the_repo_cost_map() -> None:
cost_map = json.loads((ROOT / "model_prices_and_context_window.json").read_text())
first = sync.compute_sync(cost_map, RECORDED_CATALOG, RECORDED_DOC)
second = sync.compute_sync(first.cost_map, RECORDED_CATALOG, RECORDED_DOC)
assert not second.has_changes
assert second.cost_map == first.cost_map
def test_pr_body_lists_every_section_and_the_skipped_types() -> None:
outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC)
body = sync.render_pr_body(outcome)
assert "### Added (26)" in body
assert "### Warnings needing a human call" in body
assert "image (29)" in body
assert "video (38)" in body

View file

@ -1086,6 +1086,7 @@ _COST_MAP_ROW_SCHEMA: Final = {
"chat",
"completion",
"container",
"decisions",
"image_edit",
"embedding",
"evaluation",

View file

@ -94,6 +94,7 @@ CONNECTION_NAMES: Final = (
"github_copilot_host",
"configurable_clientside_auth_params",
"use_xai_oauth",
"chatgpt_auth_file",
"fireworks_forward_user_id",
"aws_batch_role_arn",
"aws_bedrock_control_plane_endpoint",

View file

@ -16,8 +16,8 @@ from litellm.types.router import (
holds_secret_pointer,
reject_server_owned_wif_params,
server_owned_wif_fields_named,
server_owned_wif_fields_written,
server_owned_wif_fields_present,
server_owned_wif_fields_written,
)
from litellm.types.utils import (
CustomPricingLiteLLMParams,
@ -367,6 +367,7 @@ def test_server_owned_registry_includes_anthropic_openai_and_oauth_token_exchang
+ openai_wif_litellm_params
+ oauth_token_exchange_litellm_params
+ github_copilot_oauth_litellm_params
+ ("chatgpt_auth_file",)
)
assert set(openai_wif_litellm_params) == {
"openai_identity_provider_id",

View file

@ -30,7 +30,7 @@ Type tests are `*.test-d.ts` files run by the `types` vitest project (`npm run t
<!-- BEGIN:nextjs-agent-rules -->
## This is NOT the Next.js you know
# This is NOT the Next.js you know
This version has breaking changes — APIs, conventions, and file structure may all differ from your training data. Read the relevant guide in `node_modules/next/dist/docs/` (resolved from this file's directory; in monorepos the `next` package may not be visible from the repo root) before writing any code. Heed deprecation notices.

View file

@ -121,6 +121,7 @@ export interface AutoRouterDeployment extends AutoRouterCandidateDeployment {
/** False for config.yaml-defined deployments, which the update and delete routes refuse. */
db_model?: boolean | null;
base_model?: string | null;
blocked?: boolean | null;
mode?: string | null;
created_at?: string | null;
updated_at?: string | null;

View file

@ -106,10 +106,7 @@ export function AutoRoutersPanel({
>
<DialogHeader>
<DialogTitle>Add Auto Router</DialogTitle>
<DialogDescription>
Choose a classifier to route each request to a model. Called like any other model, so clients keep using a
single model name.
</DialogDescription>
<DialogDescription>Choose models for your auto router.</DialogDescription>
</DialogHeader>
<AddAutoRouterTab
handleOk={handleCreated}

View file

@ -1,4 +1,4 @@
import { render, screen, waitFor } from "@testing-library/react";
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { SearchToolTester } from "./SearchToolTester";
@ -271,7 +271,15 @@ describe("SearchToolTester", () => {
});
it("should handle search errors and show notification", async () => {
const error = new Error("Search failed");
const error = new Error(
JSON.stringify({
error: {
message: "User not allowed to access search tool: exa-search-2",
type: "user_search_tool_access_denied",
code: "403",
},
}),
);
vi.mocked(networking.searchToolQueryCall).mockRejectedValue(error);
const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {});
const user = userEvent.setup();
@ -281,11 +289,29 @@ describe("SearchToolTester", () => {
const searchButton = screen.getByRole("button", { name: /search/i });
await user.click(searchButton);
await waitFor(() => {
expect(toast.fromError).toHaveBeenCalledWith("Failed to query search tool");
expect(toast.fromError).toHaveBeenCalledWith(error);
});
consoleSpy.mockRestore();
});
it("should not show empty results or add history when search fails", async () => {
vi.mocked(networking.searchToolQueryCall).mockRejectedValue(new Error("Search failed"));
const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {});
const user = userEvent.setup();
render(<SearchToolTester {...defaultProps} />);
const input = screen.getByPlaceholderText("Enter your search query...");
fireEvent.change(input, { target: { value: "test query" } });
await user.click(screen.getByRole("button", { name: /search/i }));
await waitFor(() => {
expect(toast.fromError).toHaveBeenCalled();
});
expect(screen.queryByText("No results found")).not.toBeInTheDocument();
expect(screen.queryByText("test query")).not.toBeInTheDocument();
expect(screen.getByText("Test your search tool")).toBeInTheDocument();
consoleSpy.mockRestore();
});
it("should maintain search history", async () => {
const user = userEvent.setup();
render(<SearchToolTester {...defaultProps} />);

View file

@ -60,7 +60,7 @@ export const SearchToolTester: React.FC<SearchToolTesterProps> = ({ searchToolNa
setSearchHistory((prev) => [historyEntry, ...prev]);
} catch (error) {
console.error("Error querying search tool:", error);
toast.fromError("Failed to query search tool");
toast.fromError(error);
} finally {
setIsLoading(false);
}

View file

@ -124,7 +124,8 @@ describe("DashboardHeader breadcrumb", () => {
const tools = await screen.findByRole("dialog", { name: "Gateway tools" });
expect(within(tools).getByRole("button", { name: "Switch workspace, current: AI Gateway" })).toBeInTheDocument();
expect(within(tools).getByRole("link", { name: "Docs" })).toBeInTheDocument();
expect(within(tools).getByRole("button", { name: "Aria Agent" })).toBeInTheDocument();
expect(within(screen.getByRole("banner")).getByRole("button", { name: "Aria Agent" })).toBeInTheDocument();
expect(screen.getAllByRole("button", { name: "Aria Agent" })).toHaveLength(1);
});
it("closes mobile tools when switching to desktop and keeps them closed when returning", async () => {

View file

@ -73,43 +73,46 @@ export function DashboardHeader({
<CommandPaletteTrigger />
<div className="flex flex-none items-center gap-1 max-md:hidden">
{showWorkerSwitch && (
<>
<WorkerDropdown onWorkerSwitch={handleWorkerSwitch} />
<ToolbarSeparator />
</>
)}
<div className="flex shrink-0 items-center gap-1">
<LiteAdmin />
<DocsLink />
<BlogDropdown />
{!hideCommunityLinks && <CommunityEngagementButtons />}
<ToolbarSeparator />
<ThemeToggle />
{deploymentHealth}
<NotificationsBell accessToken={accessToken} />
</div>
<div className="flex shrink-0 items-center gap-1 md:hidden">
<CommandPaletteTrigger mobile />
{deploymentHealth}
<NotificationsBell accessToken={accessToken} />
<Popover open={!isDesktop && mobileToolsOpen} onOpenChange={setMobileToolsOpen}>
<PopoverTrigger render={<Button variant="ghost" size="icon" className="size-11" aria-label="More options" />}>
<Ellipsis />
</PopoverTrigger>
<PopoverContent align="end" className="max-w-[calc(100vw-2rem)]">
<PopoverTitle>Gateway tools</PopoverTitle>
<ViewSwitcher />
{showWorkerSwitch && <WorkerDropdown onWorkerSwitch={handleWorkerSwitch} />}
<div className="flex flex-wrap items-center gap-2">
<LiteAdmin />
<ThemeToggle />
<DocsLink />
<BlogDropdown />
{!hideCommunityLinks && <CommunityEngagementButtons />}
</div>
</PopoverContent>
</Popover>
<div className="flex flex-none items-center gap-1 max-md:hidden">
{showWorkerSwitch && (
<>
<WorkerDropdown onWorkerSwitch={handleWorkerSwitch} />
<ToolbarSeparator />
</>
)}
<DocsLink />
<BlogDropdown />
{!hideCommunityLinks && <CommunityEngagementButtons />}
<ToolbarSeparator />
<ThemeToggle />
{deploymentHealth}
<NotificationsBell accessToken={accessToken} />
</div>
<div className="flex shrink-0 items-center gap-1 md:hidden">
<CommandPaletteTrigger mobile />
{deploymentHealth}
<NotificationsBell accessToken={accessToken} />
<Popover open={!isDesktop && mobileToolsOpen} onOpenChange={setMobileToolsOpen}>
<PopoverTrigger
render={<Button variant="ghost" size="icon" className="size-11" aria-label="More options" />}
>
<Ellipsis />
</PopoverTrigger>
<PopoverContent align="end" className="max-w-[calc(100vw-2rem)]">
<PopoverTitle>Gateway tools</PopoverTitle>
<ViewSwitcher />
{showWorkerSwitch && <WorkerDropdown onWorkerSwitch={handleWorkerSwitch} />}
<div className="flex flex-wrap items-center gap-2">
<ThemeToggle />
<DocsLink />
<BlogDropdown />
{!hideCommunityLinks && <CommunityEngagementButtons />}
</div>
</PopoverContent>
</Popover>
</div>
</div>
</header>
);

View file

@ -0,0 +1,125 @@
import React from "react";
import { Select, SelectContent, SelectItem, SelectSeparator, SelectTrigger, SelectValue } from "@/components/ui/select";
import { type AutoRouterPreset } from "@/lib/autorouter_presets";
export type PresetAvailability =
| { kind: "available"; viaDeployments: boolean }
| { kind: "loading" }
| { kind: "unverifiable" }
| { kind: "missing_models"; models: readonly string[] };
// Every non-"available" state disables the option. Selection derives from this same function
// (see presetAvailability below), so an option a caller can click is always one that can be applied.
const presetDisabledHint = (availability: PresetAvailability): string | null => {
switch (availability.kind) {
case "available":
return null;
case "loading":
return "Checking model availability...";
case "unverifiable":
return "Cannot verify these models are available";
case "missing_models":
return `Missing: ${availability.models.join(", ")}`;
}
};
// "loading"/"unverifiable" are transient system states, not a gap specific to this preset; only a
// caller-specific missing-model reason gets the alarming red treatment.
const isPresetHintAlarming = (availability: PresetAvailability): boolean => availability.kind === "missing_models";
type Props = {
selectedPreset?: string;
presets: AutoRouterPreset[];
presetAvailability: (preset: AutoRouterPreset) => PresetAvailability;
onPresetChange: (key: string | undefined) => void;
pending: boolean;
unavailable: boolean;
onRetry: () => void;
};
export default function AutoRouterPresetSelect({
selectedPreset,
presets,
presetAvailability,
onPresetChange,
pending,
unavailable,
onRetry,
}: Props) {
const sortedPresetOptions = React.useMemo(
() =>
presets
.map((preset) => ({ preset, availability: presetAvailability(preset) }))
.sort((a, b) => Number(b.availability.kind === "available") - Number(a.availability.kind === "available")),
[presets, presetAvailability],
);
const templateItems = React.useMemo(
() => [
...sortedPresetOptions.map(({ preset }) => ({ value: preset.key, label: preset.label })),
{ value: "custom", label: "Start from scratch" },
],
[sortedPresetOptions],
);
return (
<div className="mb-4 w-full space-y-2">
<div className="space-y-2">
<Select
items={templateItems}
value={selectedPreset ?? null}
onValueChange={(presetKey: string | null) => onPresetChange(presetKey ?? undefined)}
>
<SelectTrigger aria-label="Use a template" data-testid="template-selector" className="w-full">
<SelectValue placeholder="Use a template" />
</SelectTrigger>
<SelectContent className="[&_[data-slot=select-item]>div]:min-w-0 [&_[data-slot=select-item]>div]:shrink [&_[data-slot=select-item]>div]:whitespace-normal [&_[data-slot=select-item]>div]:break-words">
{sortedPresetOptions.map(({ preset, availability: presetState }) => {
const disabledHint = presetDisabledHint(presetState);
const hintClass = isPresetHintAlarming(presetState) ? "text-destructive" : "text-muted-foreground";
const matchedHint =
presetState.kind === "available" && presetState.viaDeployments ? "Matches your deployments" : null;
return (
<SelectItem
key={preset.key}
value={preset.key}
label={preset.label}
disabled={disabledHint !== null}
title={disabledHint ?? preset.description}
>
<div>
<div className="font-medium">{preset.label}</div>
<div className="text-xs text-muted-foreground">{preset.description}</div>
<div className="text-xs text-muted-foreground">Replaces classifier and model choices</div>
{disabledHint && <div className={`text-xs mt-1 ${hintClass}`}>{disabledHint}</div>}
{matchedHint && <div className="text-xs mt-1 text-success">{matchedHint}</div>}
</div>
</SelectItem>
);
})}
<SelectSeparator className="mx-0" />
<SelectItem value="custom" label="Start from scratch">
<div>
<div className="font-medium">Start from scratch</div>
<div className="text-xs text-muted-foreground">
Clears tier models, switches to Heuristics, and resets routing settings
</div>
</div>
</SelectItem>
</SelectContent>
</Select>
<p className="text-xs text-muted-foreground">Replaces classifier and model choices</p>
{pending && <div className="text-xs mt-1 text-muted-foreground">Loading templates...</div>}
{unavailable && (
<div className="text-xs mt-1 text-destructive">
Could not load templates, so only Start from scratch is shown.{" "}
<button type="button" className="underline" onClick={onRetry}>
Retry
</button>
</div>
)}
</div>
</div>
);
}

View file

@ -0,0 +1,190 @@
import { Button } from "@/components/ui/button";
import { Label } from "@/components/ui/label";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import { ROUTER_FAMILIES, ROUTER_GOALS, type FamilyPreference, type RouterGoal, type SolverModel } from "./auto_setup";
type Props = {
customized: boolean;
hasSuggestions: boolean;
onReset: () => void;
goal: RouterGoal;
family: FamilyPreference;
models: readonly SolverModel[];
everyday?: SolverModel;
hard?: SolverModel;
waitingForTeam: boolean;
loading: boolean;
error: boolean;
notice: string | null;
onGoalChange: (goal: RouterGoal) => void;
onFamilyChange: (family: FamilyPreference) => void;
onModelChange: (role: "everyday" | "hard", model: string) => void;
onAdvanced: () => void;
onRetry: () => void;
};
const sharedFamily = (everyday?: SolverModel, hard?: SolverModel) =>
everyday?.family && everyday.family === hard?.family ? ROUTER_FAMILIES[everyday.family] : null;
const modelDescription = (model: SolverModel) =>
[model.model !== model.name ? model.model : null, model.provider].filter(Boolean).join(" · ");
export default function AutoRouterSimpleSetup({
customized,
hasSuggestions,
onReset,
goal,
family,
models,
everyday,
hard,
waitingForTeam,
loading,
error,
notice,
onGoalChange,
onFamilyChange,
onModelChange,
onAdvanced,
onRetry,
}: Props) {
const inventoryReady = !waitingForTeam && !loading;
const showPreferences = !inventoryReady || error || hasSuggestions;
const sameFamily = sharedFamily(everyday, hard);
const familyItems = [
{ value: "auto", label: sameFamily ? `Auto (${sameFamily})` : "Auto" },
...Object.entries(ROUTER_FAMILIES).map(([value, label]) => ({ value, label })),
];
return (
<div className="space-y-4">
{showPreferences && (
<>
<div className="space-y-2">
<Label>Optimize for</Label>
<div className="grid grid-cols-3 gap-2" role="group" aria-label="Routing goal">
{Object.entries(ROUTER_GOALS).map(([key, label]) => (
<Button
key={key}
type="button"
className="w-full"
variant={goal === key ? "default" : "outline"}
aria-pressed={goal === key}
onClick={() => onGoalChange(key as RouterGoal)}
>
{label}
</Button>
))}
</div>
</div>
<div className="space-y-2">
<Label htmlFor="auto-router-family">Model family</Label>
<Select
items={familyItems}
value={family}
onValueChange={(value) => {
if (value === "auto" || (value && value in ROUTER_FAMILIES)) onFamilyChange(value as FamilyPreference);
}}
>
<SelectTrigger id="auto-router-family" aria-label="Model family" className="w-full">
<SelectValue />
</SelectTrigger>
<SelectContent>
{familyItems.map(({ value, label }) => (
<SelectItem key={value} value={value}>
{label}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
</>
)}
{waitingForTeam && <p className="text-sm text-muted-foreground">Select a team to see available models.</p>}
{!waitingForTeam && loading && (
<p className="flex items-center gap-2 text-sm text-muted-foreground" role="status">
<UiLoadingSpinner className="size-4" />
Loading deployed models
</p>
)}
{inventoryReady && error && (
<p role="alert" className="text-sm text-destructive">
Could not load deployed models.{" "}
<Button type="button" variant="link" onClick={onRetry}>
Retry
</Button>
</p>
)}
{inventoryReady && !error && (
<>
<div className="space-y-2">
<div className="flex items-center justify-between gap-2">
<Label>Chosen models</Label>
{customized && hasSuggestions && (
<Button type="button" variant="link" className="h-auto p-0 text-xs" onClick={onReset}>
Use suggested models
</Button>
)}
</div>
<div className="grid gap-3 sm:grid-cols-2">
{(["everyday", "hard"] as const).map((role) => {
const chosen = role === "everyday" ? everyday : hard;
const other = role === "everyday" ? hard : everyday;
const choices = models.filter((model) => model.name !== other?.name);
const items = choices.map((model) => ({ value: model.name, label: model.name }));
const label = role === "everyday" ? "Everyday" : "Hard tasks";
return (
<div className="min-w-0 space-y-3 rounded-lg border bg-muted/20 p-4" key={role}>
<div className="space-y-1">
<Label htmlFor={`auto-router-${role}`}>{label}</Label>
<p className="text-xs text-muted-foreground">
{role === "everyday" ? "Your cheaper model" : "Your stronger model"}
</p>
</div>
<Select
items={items}
value={chosen?.name ?? null}
onValueChange={(value) => {
if (value) onModelChange(role, value);
}}
>
<SelectTrigger id={`auto-router-${role}`} aria-label={`${label} model`} className="w-full">
<SelectValue placeholder="Select a deployed model" />
</SelectTrigger>
<SelectContent>
{choices.map((model) => (
<SelectItem key={model.name} value={model.name} label={model.name}>
<div>
<div>{model.name}</div>
<div className="text-xs text-muted-foreground">{modelDescription(model)}</div>
</div>
</SelectItem>
))}
</SelectContent>
</Select>
{chosen && <p className="break-words text-xs text-muted-foreground">{modelDescription(chosen)}</p>}
</div>
);
})}
</div>
</div>
{(!everyday || !hard) && (
<p role="status" className="text-sm text-muted-foreground">
{models.length < 2
? "Two deployed models are needed. Add a model or use Advanced configuration."
: "Choose two models to get started."}
</p>
)}
{notice && !customized && (
<p role="status" className="text-xs text-muted-foreground">
{notice}
</p>
)}
</>
)}
<Button type="button" variant="link" className="px-0" onClick={onAdvanced}>
Advanced configuration
</Button>
</div>
);
}

View file

@ -13,7 +13,7 @@ import { FormField } from "@/components/shared/form/FormField";
import { Button } from "@/components/ui/button";
import { Card, CardContent } from "@/components/ui/card";
import { Input } from "@/components/ui/input";
import { Select, SelectContent, SelectItem, SelectSeparator, SelectTrigger, SelectValue } from "@/components/ui/select";
import AutoRouterPresetSelect, { type PresetAvailability } from "./AutoRouterPresetSelect";
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import { useZodForm } from "@/lib/forms/useZodForm";
@ -79,7 +79,16 @@ import {
} from "@/lib/autorouter_presets";
import { useAutoRouterPresets } from "@/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets";
import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
import { buildAutomaticRouterConfig, buildPreferredTierModels, prefillEmptyTiers } from "./auto_setup";
import AutoRouterSimpleSetup from "./AutoRouterSimpleSetup";
import {
deployedRouterModels,
nextRouterName,
simpleSetupError,
recommendSolverPair,
solverPairConfig,
type RouterGoal,
type FamilyPreference,
} from "./auto_setup";
interface AddAutoRouterTabProps {
handleOk: () => void;
@ -95,31 +104,6 @@ interface AddAutoRouterTabProps {
teams?: Team[] | null;
}
type PresetAvailability =
| { kind: "available"; viaDeployments: boolean }
| { kind: "loading" }
| { kind: "unverifiable" }
| { kind: "missing_models"; models: readonly string[] };
// Every non-"available" state disables the option. Selection derives from this same function
// (see presetAvailability below), so an option a caller can click is always one that can be applied.
const presetDisabledHint = (availability: PresetAvailability): string | null => {
switch (availability.kind) {
case "available":
return null;
case "loading":
return "Checking model availability...";
case "unverifiable":
return "Cannot verify these models are available";
case "missing_models":
return `Missing: ${availability.models.join(", ")}`;
}
};
// "loading"/"unverifiable" are transient system states, not a gap specific to this preset; only a
// caller-specific missing-model reason gets the alarming red treatment.
const isPresetHintAlarming = (availability: PresetAvailability): boolean => availability.kind === "missing_models";
const NO_PRESETS: AutoRouterPreset[] = [];
// Why the submit is unavailable, or null when it is available. The button reads this to disable
@ -168,7 +152,7 @@ const autoRouterSchema = (requiresTeamScope: boolean) =>
type AddAutoRouterFormValues = z.infer<ReturnType<typeof autoRouterSchema>>;
const EMPTY_FORM_VALUES: AddAutoRouterFormValues = {
auto_router_name: "",
auto_router_name: "auto-router",
team_id: null,
model_access_group: undefined,
};
@ -210,11 +194,18 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
});
const [modelAccessGroups, setModelAccessGroups] = useState<string[]>([]);
const [complexityRouterConfig, setComplexityRouterConfig] = useState<ComplexityRouterConfigValue>({
const [advancedConfig, setComplexityRouterConfig] = useState<ComplexityRouterConfigValue>({
tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] },
classifier_type: "heuristic",
classification_mode: "user_turn",
});
const [advanced, setAdvanced] = useState(false);
const [prefillAttempted, setPrefillAttempted] = useState(false);
const [goal, setGoal] = useState<RouterGoal>("balanced");
const [family, setFamily] = useState<FamilyPreference>("auto");
const [solverOverrides, setSolverOverrides] = useState<{ everyday?: string; hard?: string }>({});
const [customTechnicalKeywords, setCustomTechnicalKeywords] = useState<string[]>([]);
const [keywordTierRules, setKeywordTierRules] = useState<KeywordTierRule[]>([]);
const [semanticMatchingEnabled, setSemanticMatchingEnabled] = useState<boolean>(false);
@ -250,17 +241,30 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
isError: modelsError,
refetch: refetchModels,
} = useQuery({
queryKey: ["availableModels", "autoRouter", accessToken, ...(isMemberManaged ? [watchedTeamId] : [])],
queryKey: ["availableModels", "autoRouter", accessToken, ...(requiresTeamScope ? [watchedTeamId] : [])],
queryFn: () =>
isMemberManaged ? fetchAutoRouterModels(accessToken, watchedTeamId) : fetchAvailableModels(accessToken),
enabled: Boolean(accessToken && (!isMemberManaged || watchedTeamId)),
requiresTeamScope ? fetchAutoRouterModels(accessToken, watchedTeamId) : fetchAvailableModels(accessToken),
enabled: Boolean(accessToken && (!requiresTeamScope || watchedTeamId)),
});
const { data: deployments, isLoading: deploymentsLoading } = useQuery({
const {
data: deployments,
isLoading: deploymentsLoading,
isError: deploymentsError,
refetch: refetchDeployments,
} = useQuery({
queryKey: autoRouterListKey(userId ?? "", userRole),
queryFn: () => fetchAllModelDeployments(accessToken, userId ?? "", userRole),
enabled: Boolean(accessToken),
});
const modelsLoading = groupsLoading || deploymentsLoading;
useEffect(() => {
if (!deployments || form.getFieldState("auto_router_name").isDirty) return;
form.setValue(
"auto_router_name",
nextRouterName(deployments.flatMap((row) => (row.model_name ? [row.model_name] : []))),
);
}, [deployments, form]);
const modelInfo = React.useMemo(() => data ?? [], [data]);
const {
data: presetsData,
@ -269,12 +273,12 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
refetch: refetchPresets,
} = useAutoRouterPresets();
const presets = presetsData ?? NO_PRESETS;
const automaticSetupLoading = modelsLoading || presetsPending;
const presetsUnavailable = presetsError && presetsData === undefined;
// react-query keeps the last successful list around when a later refetch fails, so isError alone
// can't tell "never loaded" apart from "loaded, then a background refetch errored" - only the
// former leaves us with nothing trustworthy to verify a preset's models against.
const modelsUnverifiable = modelsError && data === undefined;
const inventoryUnavailable = modelsUnverifiable || (deploymentsError && deployments === undefined);
const isAdmin = all_admin_roles.includes(userRole);
@ -294,14 +298,33 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
),
[modelInfo],
);
const preferredTierModels = React.useMemo(
() => buildPreferredTierModels(presets, availability),
[presets, availability],
const solverModels = React.useMemo(
() => deployedRouterModels(modelInfo, deployments ?? []),
[modelInfo, deployments],
);
const automaticRouterConfig = React.useMemo(
() => buildAutomaticRouterConfig(modelInfo, deployments ?? [], preferredTierModels),
[modelInfo, deployments, preferredTierModels],
const recommended = React.useMemo(
() => recommendSolverPair(solverModels, goal, family),
[solverModels, goal, family],
);
const everyday = solverOverrides.everyday
? solverModels.find((model) => model.name === solverOverrides.everyday)
: recommended?.everyday;
const hard = solverOverrides.hard
? solverModels.find((model) => model.name === solverOverrides.hard)
: recommended?.hard;
const simpleConfig = everyday && hard ? solverPairConfig(everyday.name, hard.name) : null;
const complexityRouterConfig = advanced
? advancedConfig
: simpleConfig ?? {
classifier_type: "heuristic" as const,
classification_mode: "user_turn" as const,
tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] },
};
const openAdvanced = () => {
setComplexityRouterConfig(complexityRouterConfig);
setPrefillAttempted(simpleConfig !== null);
setAdvanced(true);
};
// A preset's models can only be trusted against a successfully loaded list. Selection and the
// greyed-out state derive from this one function, so a preset that cannot be selected can never
@ -322,22 +345,6 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
[modelsLoading, modelsUnverifiable, availability, groupsOnlyAvailability],
);
const sortedPresetOptions = React.useMemo(
() =>
presets
.map((preset) => ({ preset, availability: presetAvailability(preset) }))
.sort((a, b) => Number(b.availability.kind === "available") - Number(a.availability.kind === "available")),
[presets, presetAvailability],
);
const templateItems = React.useMemo(
() => [
...sortedPresetOptions.map(({ preset }) => ({ value: preset.key, label: preset.label })),
{ value: "custom", label: "Start from scratch" },
],
[sortedPresetOptions],
);
const applyPrefill = (prefill: PresetPrefill) => {
setEditingTiers(false);
setComplexityRouterConfig(prefill.complexityRouterConfig);
@ -349,16 +356,20 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
setEscalationKeywords(prefill.escalationKeywords);
};
const [prefillAttempted, setPrefillAttempted] = useState(false);
const automaticPrefill = automaticSetupLoading || modelsUnverifiable ? null : automaticRouterConfig;
if (automaticPrefill && !prefillAttempted) {
if (advanced && !prefillAttempted && !modelsLoading && simpleConfig) {
setPrefillAttempted(true);
setComplexityRouterConfig((current) => prefillEmptyTiers(current, automaticPrefill));
setComplexityRouterConfig((current) =>
!current.custom_tier_set &&
!isForecastClassifier(current.classifier_type) &&
activeTierRows(current).every((row) => row.models.length === 0)
? { ...current, tiers: simpleConfig.tiers }
: current,
);
}
const handlePresetChange = (presetKey: string | undefined) => {
setPrefillAttempted(true);
if (!presetKey || presetKey === "custom") {
setPrefillAttempted(true);
setSelectedPreset(presetKey);
applyPrefill(buildEmptyPrefill());
@ -372,7 +383,6 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
const presetState = presetAvailability(preset);
if (presetState.kind !== "available") return;
setPrefillAttempted(true);
setSelectedPreset(presetKey);
applyPrefill(buildPresetPrefill(preset.complexity_router_config, availability));
};
@ -413,7 +423,11 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
},
!requiresTeamScope || Boolean(watchedTeamId),
);
const saveBlockedReason = submitBlockedReason ?? routerAvailability.saveBlockedReason;
const simpleBlockedReason = advanced ? null : simpleSetupError(modelsLoading, inventoryUnavailable, everyday, hard);
const saveBlockedReason =
!advanced && requiresTeamScope && !watchedTeamId
? "Select a team to see available models."
: simpleBlockedReason ?? submitBlockedReason ?? routerAvailability.saveBlockedReason;
const handleClassifierChange = (config: ComplexityRouterConfigValue) => {
setSelectedPreset(undefined);
setComplexityRouterConfig(config);
@ -540,71 +554,20 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
setIsTestModalVisible(true);
};
const tierSetup = (
<div className="mb-4 w-full space-y-2">
<div className="space-y-2">
<Select
items={templateItems}
value={selectedPreset ?? null}
onValueChange={(presetKey: string | null) => handlePresetChange(presetKey ?? undefined)}
>
<SelectTrigger aria-label="Use a template" data-testid="template-selector" className="w-full">
<SelectValue placeholder="Use a template" />
</SelectTrigger>
<SelectContent className="[&_[data-slot=select-item]>div]:min-w-0 [&_[data-slot=select-item]>div]:shrink [&_[data-slot=select-item]>div]:whitespace-normal [&_[data-slot=select-item]>div]:break-words">
{sortedPresetOptions.map(({ preset, availability: presetState }) => {
const disabledHint = presetDisabledHint(presetState);
const hintClass = isPresetHintAlarming(presetState) ? "text-destructive" : "text-muted-foreground";
const matchedHint =
presetState.kind === "available" && presetState.viaDeployments ? "Matches your deployments" : null;
return (
<SelectItem
key={preset.key}
value={preset.key}
label={preset.label}
disabled={disabledHint !== null}
title={disabledHint ?? preset.description}
>
<div>
<div className="font-medium">{preset.label}</div>
<div className="text-xs text-muted-foreground">{preset.description}</div>
<div className="text-xs text-muted-foreground">Replaces classifier and model choices</div>
{disabledHint && <div className={`text-xs mt-1 ${hintClass}`}>{disabledHint}</div>}
{matchedHint && <div className="text-xs mt-1 text-success">{matchedHint}</div>}
</div>
</SelectItem>
);
})}
<SelectSeparator className="mx-0" />
<SelectItem value="custom" label="Start from scratch">
<div>
<div className="font-medium">Start from scratch</div>
<div className="text-xs text-muted-foreground">
Clears tier models, switches to Heuristics, and resets routing settings
</div>
</div>
</SelectItem>
</SelectContent>
</Select>
<p className="text-xs text-muted-foreground">Replaces classifier and model choices</p>
{presetsPending && <div className="text-xs mt-1 text-muted-foreground">Loading templates...</div>}
{presetsUnavailable && (
<div className="text-xs mt-1 text-destructive">
Could not load templates, so only Start from scratch is shown.{" "}
<button type="button" className="underline" onClick={() => void refetchPresets()}>
Retry
</button>
</div>
)}
</div>
</div>
);
const configurationForm = (
<ComplexityRouterConfig
creationFlow
tierSetup={tierSetup}
tierSetup={
<AutoRouterPresetSelect
selectedPreset={selectedPreset}
presets={presets}
presetAvailability={presetAvailability}
onPresetChange={handlePresetChange}
pending={presetsPending}
unavailable={presetsUnavailable}
onRetry={() => void refetchPresets()}
/>
}
editingTiers={editingTiers}
onEditingTiersChange={setEditingTiers}
modelInfo={modelInfo}
@ -646,8 +609,8 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
)}
</FormField>
</div>
<div className="mb-6 border-b pb-6">
{isAdmin && (
<div className={advanced ? "mb-6 border-b pb-6" : ""}>
{isAdmin && advanced && (
<div>
<FormField
control={form.control}
@ -691,91 +654,138 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
</FormField>
)}
<AutoRouterClassifierTabs
value={complexityRouterConfig}
onChange={handleClassifierChange}
decisionModel={
<DecisionModelSelect
value={complexityRouterConfig}
onChange={setComplexityRouterConfig}
modelInfo={modelInfo}
deployments={deployments}
/>
}
judgeModel={
<ClassifierPrimarySettings
section="judge"
value={complexityRouterConfig}
onChange={setComplexityRouterConfig}
modelOptions={modelInfo
.filter((model) => model.mode !== "embedding")
.map((model) => ({ value: model.model_group, label: model.model_group }))}
showValidationErrors={showValidationErrors}
/>
}
>
<FieldGroup>
{modelsUnverifiable && (
<p className="text-sm text-destructive">
Could not load available models.{" "}
<button type="button" className="underline" onClick={() => void refetchModels()}>
Retry
</button>
</p>
)}
{configurationForm}
<div className="flex flex-wrap items-center justify-between gap-3">
<Tooltip>
<TooltipTrigger
render={
<a
href="https://github.com/BerriAI/litellm/issues"
className="text-sm text-primary underline-offset-4 hover:underline"
>
Need Help?
</a>
}
/>
<TooltipContent>Get help on our github</TooltipContent>
</Tooltip>
<div className="flex flex-wrap gap-2">
<BlockedReasonTooltip reason={submitBlockedReason}>
<Button
type="button"
variant="outline"
data-testid="auto-router-test-routing-btn"
disabled={submitBlockedReason !== null || isSubmitting}
onClick={() => setIsRoutingTestVisible(true)}
>
Test Routing
</Button>
</BlockedReasonTooltip>
<Button
type="button"
variant="outline"
data-testid="auto-router-test-connect-btn"
onClick={handleTestConnection}
disabled={isTestingConnection}
{!advanced && (
<AutoRouterSimpleSetup
customized={Boolean(solverOverrides.everyday || solverOverrides.hard)}
hasSuggestions={Boolean(recommended || recommendSolverPair(solverModels, goal, "auto"))}
onReset={() => setSolverOverrides({})}
goal={goal}
family={family}
models={solverModels}
everyday={everyday}
hard={hard}
waitingForTeam={requiresTeamScope && !watchedTeamId}
loading={modelsLoading}
error={inventoryUnavailable}
notice={
Object.keys(solverOverrides).length > 0
? "Your model choices are kept."
: recommended?.notice ?? null
}
onGoalChange={setGoal}
onFamilyChange={(next) => {
setFamily(next);
setSolverOverrides({});
}}
onModelChange={(role, model) =>
setSolverOverrides({ everyday: everyday?.name, hard: hard?.name, [role]: model })
}
onAdvanced={openAdvanced}
onRetry={() => {
void refetchModels();
void refetchDeployments();
}}
/>
)}
{advanced && (
<AutoRouterClassifierTabs
value={complexityRouterConfig}
onChange={handleClassifierChange}
decisionModel={
<DecisionModelSelect
value={complexityRouterConfig}
onChange={setComplexityRouterConfig}
modelInfo={modelInfo}
deployments={deployments}
/>
}
judgeModel={
<ClassifierPrimarySettings
section="judge"
value={complexityRouterConfig}
onChange={setComplexityRouterConfig}
modelOptions={modelInfo
.filter((model) => model.mode !== "embedding")
.map((model) => ({ value: model.model_group, label: model.model_group }))}
showValidationErrors={showValidationErrors}
/>
}
>
<FieldGroup>
{modelsUnverifiable && (
<p className="text-sm text-destructive">
Could not load available models.{" "}
<button type="button" className="underline" onClick={() => void refetchModels()}>
Retry
</button>
</p>
)}
{configurationForm}
</FieldGroup>
</AutoRouterClassifierTabs>
)}
{!advanced && routerAvailability.isError && (
<p className="mt-3 text-sm text-destructive" role="alert">
Could not check availability.{" "}
<Button type="button" variant="link" onClick={() => void routerAvailability.refetch()}>
Retry
</Button>
</p>
)}
{!advanced && routerAvailability.data?.error && (
<p className="mt-3 text-sm text-destructive" role="alert">
{routerAvailability.data.error}
</p>
)}
<div className="mt-5 flex flex-wrap items-center justify-between gap-3">
<Tooltip>
<TooltipTrigger
render={
<a
href="https://github.com/BerriAI/litellm/issues"
className="text-sm text-primary underline-offset-4 hover:underline"
>
{isTestingConnection && <UiLoadingSpinner className="size-4" />}
Test Connection
</Button>
<BlockedReasonTooltip reason={saveBlockedReason}>
<Button
type="button"
disabled={saveBlockedReason !== null || isSubmitting}
onClick={() => {
void handleAutoRouterSubmit();
}}
>
Add Auto Router
</Button>
</BlockedReasonTooltip>
</div>
</div>
</FieldGroup>
</AutoRouterClassifierTabs>
Need Help?
</a>
}
/>
<TooltipContent>Get help on our github</TooltipContent>
</Tooltip>
<div className="flex flex-wrap gap-2">
<BlockedReasonTooltip reason={submitBlockedReason}>
<Button
type="button"
variant="outline"
data-testid="auto-router-test-routing-btn"
disabled={submitBlockedReason !== null || isSubmitting}
onClick={() => setIsRoutingTestVisible(true)}
>
Test Routing
</Button>
</BlockedReasonTooltip>
<Button
type="button"
variant="outline"
data-testid="auto-router-test-connect-btn"
onClick={handleTestConnection}
disabled={isTestingConnection}
>
{isTestingConnection && <UiLoadingSpinner className="size-4" />}
Test Connection
</Button>
<BlockedReasonTooltip reason={saveBlockedReason}>
<Button
type="button"
disabled={saveBlockedReason !== null || isSubmitting}
onClick={() => {
void handleAutoRouterSubmit();
}}
>
Add Auto Router
</Button>
</BlockedReasonTooltip>
</div>
</div>
</form>
</CardContent>
</Card>

View file

@ -1,274 +1,325 @@
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
import { readFileSync } from "node:fs";
import { describe, expect, it } from "vitest";
import type { AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels";
import { buildModelAvailability, deploymentRefsFromModelInfo } from "@/lib/autorouter_presets";
import {
buildAutomaticRouterConfig,
buildPreferredTierModels,
prefillEmptyTiers,
type PreferredTierModels,
deployedRouterModels,
recommendSolverPair,
solverPairConfig,
nextRouterName,
simpleSetupError,
} from "./auto_setup";
const models = (...names: string[]) => names.map((model_group) => ({ model_group, mode: "chat" }));
const reasoningModel = (model_group: string, supported_reasoning_efforts: string[]) => ({
model_group,
mode: "chat",
supports_reasoning: true,
supported_reasoning_efforts,
});
const deployment = (model_name: string, model = model_name): AutoRouterDeployment => ({
model_name,
const deployment = (name: string, model: string, extra: Partial<AutoRouterDeployment> = {}): AutoRouterDeployment => ({
model_name: name,
litellm_params: { model },
...extra,
});
const tierModels = (config: ReturnType<typeof buildAutomaticRouterConfig>) =>
config && Object.values(config.tiers).map((tier) => (typeof tier === "string" ? tier : tier[0]));
const resolve = (...rows: AutoRouterDeployment[]) =>
deployedRouterModels(
rows.map((row) => ({ model_group: row.model_name!, mode: "chat" })),
rows,
);
const names = (pair: ReturnType<typeof recommendSolverPair>) => pair && [pair.everyday.name, pair.hard.name];
const ladder = () =>
resolve(
deployment("small", "anthropic/claude-haiku-5-5"),
deployment("medium", "anthropic/claude-sonnet-5-5"),
deployment("large", "anthropic/claude-opus-5-5"),
);
describe("buildPreferredTierModels", () => {
it("recognizes curated models that are not in a preset", () => {
const available = ["gpt-5.6-luna", "claude-sonnet-5", "grok-4.6", "deepseek-v4-pro"];
const availability = buildModelAvailability(available, []);
const preferred = buildPreferredTierModels([], availability);
const expected: PreferredTierModels = {
SIMPLE: ["gpt-5.6-luna"],
MEDIUM: ["claude-sonnet-5"],
COMPLEX: ["deepseek-v4-pro", "grok-4.6"],
REASONING: ["deepseek-v4-pro", "grok-4.6"],
};
const catalog: Record<string, { mode?: string; litellm_provider?: string }> = JSON.parse(
readFileSync(new URL("../../../../../model_prices_and_context_window.json", import.meta.url), "utf8"),
);
const catalogOpusModels = Object.entries(catalog).filter(
([id, metadata]) => metadata.mode === "chat" && /claude-opus-\d/.test(id) && !/-(?:fast|think)$/.test(id),
);
expect(preferred).toEqual(expected);
});
it("prefers the current model ladder over older preset entries", () => {
const availability = buildModelAvailability(["gpt-5.6-luna", "gpt-4o-mini"], []);
const preferred = buildPreferredTierModels(
[
{
key: "old",
label: "Old",
description: "Old model",
complexity_router_config: {
tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [], REASONING: [] },
classifier_type: "heuristic_v2",
},
},
],
availability,
);
expect(preferred.SIMPLE).toEqual(["gpt-5.6-luna", "gpt-4o-mini"]);
});
});
describe("buildAutomaticRouterConfig", () => {
it("selects native Terra and Sol groups with their reasoning settings, retaining cloud fallback", () => {
const modelNames = ["gpt-5.6-terra", "gpt-5.6-sol"];
const deployments = modelNames.flatMap((model) => [
deployment(model, `azure/${model}`),
deployment(`z-native-${model}`, `openai/${model}`),
]);
const available = deployments.map(({ model_name }) => reasoningModel(model_name!, ["none", "high"]));
const availability = buildModelAvailability(
available.map(({ model_group }) => model_group),
deploymentRefsFromModelInfo(deployments),
);
const preferred = buildPreferredTierModels([], availability);
const config = buildAutomaticRouterConfig(available, deployments, preferred);
expect(tierModels(config)).toEqual([
"z-native-gpt-5.6-terra",
"z-native-gpt-5.6-terra",
"z-native-gpt-5.6-sol",
"z-native-gpt-5.6-sol",
]);
expect(config?.tier_model_params).toEqual({
REASONING: { "z-native-gpt-5.6-sol": { reasoning_effort: "high" } },
});
const cloudOnly = models(...modelNames);
const cloudAvailability = buildModelAvailability(modelNames, deploymentRefsFromModelInfo(deployments));
const cloudPreferred = buildPreferredTierModels([], cloudAvailability);
expect(tierModels(buildAutomaticRouterConfig(cloudOnly, deployments, cloudPreferred))).toEqual([
"gpt-5.6-terra",
"gpt-5.6-terra",
"gpt-5.6-sol",
"gpt-5.6-sol",
]);
});
it("selects one preferred model for each tier", () => {
const preferred: PreferredTierModels = {
SIMPLE: ["simple"],
MEDIUM: ["medium"],
COMPLEX: ["complex"],
REASONING: ["reasoning"],
};
expect(
tierModels(buildAutomaticRouterConfig(models("simple", "medium", "complex", "reasoning"), [], preferred)),
).toEqual(["simple", "medium", "complex", "reasoning"]);
});
it.each([
{
provider: "OpenAI",
available: ["gpt-5.6-luna", "gpt-5.6-terra", "gpt-6-astra"],
expected: ["gpt-5.6-luna", "gpt-5.6-terra", "gpt-6-astra", "gpt-6-astra"],
supportedEfforts: ["low", "medium", "high", "xhigh", "max"],
effort: "max",
},
{
provider: "Anthropic",
available: ["claude-haiku-4-5", "claude-sonnet-5", "claude-opus-5"],
expected: ["claude-haiku-4-5", "claude-sonnet-5", "claude-opus-5", "claude-opus-5"],
supportedEfforts: ["low", "medium", "high", "max"],
effort: "max",
},
{
provider: "Google",
available: ["gemini-3.5-flash-lite", "gemini-3.8-flash", "gemini-3.1-pro-preview"],
expected: ["gemini-3.5-flash-lite", "gemini-3.8-flash", "gemini-3.1-pro-preview", "gemini-3.1-pro-preview"],
supportedEfforts: ["low", "medium", "high"],
effort: "high",
},
{
provider: "DeepSeek",
available: ["deepseek-v4-flash", "deepseek-v4-pro"],
expected: ["deepseek-v4-flash", "deepseek-v4-flash", "deepseek-v4-pro", "deepseek-v4-pro"],
supportedEfforts: ["none", "high"],
effort: "high",
},
{
provider: "xAI",
available: ["grok-4.6"],
expected: ["grok-4.6", "grok-4.6", "grok-4.6", "grok-4.6"],
supportedEfforts: ["low", "medium", "high", "xhigh"],
effort: "xhigh",
},
])(
"uses the current $provider ladder and strongest advertised reasoning effort",
({ available, expected, supportedEfforts, effort }) => {
const availability = buildModelAvailability(available, []);
const preferred = buildPreferredTierModels([], availability);
const modelInfo = models(...available).map((model) =>
model.model_group === expected[3] ? reasoningModel(model.model_group, supportedEfforts) : model,
);
const config = buildAutomaticRouterConfig(modelInfo, [], preferred);
expect(tierModels(config)).toEqual(expected);
expect(config?.tier_model_params).toEqual({
REASONING: { [expected[3]]: { reasoning_effort: effort } },
});
describe("deployedRouterModels", () => {
it.each(catalogOpusModels)(
"recognizes the catalog's standard Opus deployment %s behind a custom alias",
(id, metadata) => {
const row = {
model_name: "custom-deployment",
litellm_params: { model: id, custom_llm_provider: metadata.litellm_provider },
};
const models = deployedRouterModels([{ model_group: "custom-deployment", mode: "chat" }], [row]);
expect(models).toEqual([expect.objectContaining({ name: "custom-deployment", family: "claude", line: "opus" })]);
},
);
it("never exceeds the selected model group's advertised reasoning efforts", () => {
const available = ["gpt-5.6-luna", "gpt-5.6-terra", "gpt-5.6-sol"];
const availability = buildModelAvailability(available, []);
const preferred = buildPreferredTierModels([], availability);
const modelInfo = [
...models("gpt-5.6-luna", "gpt-5.6-terra"),
reasoningModel("gpt-5.6-sol", ["none", "low", "medium", "high", "xhigh"]),
];
const config = buildAutomaticRouterConfig(modelInfo, [], preferred);
expect(config?.tier_model_params).toEqual({
REASONING: { "gpt-5.6-sol": { reasoning_effort: "xhigh" } },
});
it("resolves underlying identities independently of deployment aliases and serving providers", () => {
const models = resolve(
deployment("everyday", "bedrock/us.openai.gpt-6.1-sol"),
deployment("hard", "bedrock_converse/global.openai.gpt-6-astra"),
);
const expectedEveryday = { name: "everyday", family: "openai", line: "sol", version: [6, 1], provider: "Bedrock" };
const expectedHard = { name: "hard", family: "openai", line: "astra", version: [6], provider: "Bedrock" };
expect(models).toEqual(
expect.arrayContaining([expect.objectContaining(expectedEveryday), expect.objectContaining(expectedHard)]),
);
expect(names(recommendSolverPair(models, "quality", "openai"))).toEqual(["everyday", "hard"]);
});
it("leaves reasoning effort unset when the proxy does not report supported values", () => {
const available = ["grok-4.6"];
const availability = buildModelAvailability(available, []);
const preferred = buildPreferredTierModels([], availability);
const config = buildAutomaticRouterConfig(models(...available), [], preferred);
expect(config?.tier_model_params).toBeUndefined();
});
it("reuses the closest available tier when a tier has no match", () => {
const preferred: PreferredTierModels = {
SIMPLE: ["simple"],
MEDIUM: [],
COMPLEX: ["complex"],
REASONING: [],
};
expect(tierModels(buildAutomaticRouterConfig(models("simple", "complex"), [], preferred))).toEqual([
"simple",
"simple",
"complex",
"complex",
it("recognizes provider-specific Claude names and a declared base model for opaque deployments", () => {
const models = resolve(
deployment("small", "vertex_ai/claude-haiku-4-5@20251001"),
deployment("large", "bedrock/arn:example:inference-profile/us.anthropic.claude-opus-5-5-v1:0"),
deployment("middle", "azure/private-alias", {
litellm_params: { model: "azure/private-alias" },
model_info: { base_model: "gpt-6.1-sol" },
}),
);
expect(models.map(({ name, family, provider }) => ({ name, family, provider }))).toEqual([
{ name: "large", family: "claude", provider: "Bedrock" },
{ name: "middle", family: "openai", provider: "Azure" },
{ name: "small", family: "claude", provider: "Vertex AI" },
]);
});
it("returns null when none of the available models are recommended", () => {
const preferred: PreferredTierModels = {
SIMPLE: ["missing-simple"],
MEDIUM: ["missing-medium"],
COMPLEX: ["missing-complex"],
REASONING: ["missing-reasoning"],
};
expect(buildAutomaticRouterConfig(models("unknown-model"), [], preferred)).toBeNull();
it("uses only caller-visible chat models and excludes blocked and recursive targets", () => {
const rows = [
deployment("hidden", "anthropic/claude-opus-6"),
deployment("blocked", "anthropic/claude-opus-5", { model_info: { blocked: true } }),
deployment("embedding", "anthropic/claude-opus-5"),
deployment("recursive", "auto_router/complexity_router"),
deployment("ready", "anthropic/claude-sonnet-5"),
];
expect(
deployedRouterModels(
[
{ model_group: "blocked", mode: "chat" },
{ model_group: "embedding", mode: "embedding" },
{ model_group: "recursive" },
{ model_group: "ready", mode: "chat" },
],
rows,
).map((model) => model.name),
).toEqual(["ready"]);
});
it("ignores non-chat models and existing auto routers", () => {
const preferred: PreferredTierModels = {
SIMPLE: ["gpt-4o-mini", "smart-router"],
MEDIUM: [],
COMPLEX: [],
REASONING: [],
};
const available = [
{ model_group: "gpt-4o-mini", mode: "chat" },
{ model_group: "image-model", mode: "image_generation" },
{ model_group: "smart-router", mode: "chat" },
it("does not claim a mixed-version or unknown model group is a known version", () => {
const rows = [
deployment("mixed", "anthropic/claude-opus-5"),
deployment("mixed", "anthropic/claude-opus-5-5"),
deployment("claude-sonnet-5", "openai/custom-model"),
deployment("wildcard", "openai/*"),
];
const models = deployedRouterModels(
[
{ model_group: "mixed", mode: "chat" },
{ model_group: "claude-sonnet-5", mode: "chat" },
{ model_group: "wildcard", mode: "chat" },
],
rows,
);
expect(models.map((model) => model.name)).toEqual(["claude-sonnet-5", "mixed"]);
expect(models.every((model) => model.family === undefined)).toBe(true);
expect(recommendSolverPair(models, "balanced", "auto")).toBeNull();
});
it("keeps homogeneous replicas and ignores blocked replicas when resolving identity", () => {
const rows = [
deployment("group", "anthropic/claude-opus-5-5"),
deployment("group", "bedrock/anthropic.claude-opus-5-5-v1:0"),
deployment("group", "anthropic/claude-opus-5", { model_info: { blocked: true } }),
];
const models = deployedRouterModels([{ model_group: "group" }], rows);
expect(models).toHaveLength(1);
expect(models[0]).toMatchObject({ name: "group", version: [5, 5], provider: "Anthropic, Bedrock" });
});
it("keeps arbitrary canonical names available without pretending an unversioned alias is latest", () => {
expect(
tierModels(
buildAutomaticRouterConfig(available, [deployment("smart-router", "auto_router/complexity_router")], preferred),
deployedRouterModels([{ model_group: "claude-opus-8-12" }, { model_group: "claude-opus-latest" }], []).map(
(model) => model.name,
),
).toEqual(["gpt-4o-mini", "gpt-4o-mini", "gpt-4o-mini", "gpt-4o-mini"]);
).toEqual(["claude-opus-8-12", "claude-opus-latest"]);
});
});
describe("prefillEmptyTiers", () => {
const empty: ComplexityRouterConfigValue = {
classifier_type: "jev",
jev_classifier_config: { model: "jev-custom", timeout_ms: 4200 },
classification_mode: "user_turn",
tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] },
};
const automatic: ComplexityRouterConfigValue = {
classifier_type: "heuristic_v2",
tiers: { SIMPLE: ["fast"], MEDIUM: ["mid"], COMPLEX: ["strong"], REASONING: ["strong"] },
tier_model_params: { REASONING: { strong: { reasoning_effort: "high" } } },
};
it("fills tiers while preserving a classifier configured before models load", () => {
expect(prefillEmptyTiers(empty, automatic)).toEqual({
...empty,
tiers: automatic.tiers,
tier_model_params: automatic.tier_model_params,
describe("recommendSolverPair", () => {
it.each([
["cheapest", ["small", "medium"]],
["balanced", ["small", "large"]],
["quality", ["medium", "large"]],
] as const)("uses the requested %s pair", (goal, expected) => {
expect(names(recommendSolverPair(ladder(), goal, "claude"))).toEqual(expected);
});
it("selects future deployed releases numerically within each line, without a version allowlist", () => {
const models = resolve(
deployment("old-small", "anthropic/claude-haiku-8-9"),
deployment("new-small", "bedrock/anthropic.claude-haiku-8-10-v1:0"),
deployment("old-large", "anthropic/claude-opus-8-9"),
deployment("new-large", "anthropic/claude-opus-9"),
);
expect(names(recommendSolverPair(models, "balanced", "claude"))).toEqual(["new-small", "new-large"]);
});
it("orders snapshots within a version without confusing a snapshot with the next version", () => {
const models = resolve(
deployment("old", "anthropic/claude-haiku-4-5-20251001"),
deployment("snapshot", "anthropic/claude-haiku-4-5-20251101"),
deployment("next", "anthropic/claude-haiku-4-6"),
deployment("large", "anthropic/claude-opus-5"),
);
expect(names(recommendSolverPair(models, "balanced", "claude"))).toEqual(["next", "large"]);
expect(
names(
recommendSolverPair(
models.filter((model) => model.name !== "next"),
"balanced",
"claude",
),
),
).toEqual(["snapshot", "large"]);
});
it("prefers rolling aliases over snapshots of the same release", () => {
const models = resolve(
deployment("pin", "anthropic/claude-haiku-4-5-20251101"),
deployment("alias", "anthropic/claude-haiku-4-5"),
deployment("hard", "anthropic/claude-opus-5"),
deployment("luna-pin", "openai/gpt-6-luna-2026-01-01"),
deployment("luna", "openai/gpt-6-luna"),
deployment("astra", "openai/gpt-6-astra"),
);
expect(names(recommendSolverPair(models, "balanced", "claude"))).toEqual(["alias", "hard"]);
expect(names(recommendSolverPair(models, "balanced", "openai"))).toEqual(["luna", "astra"]);
});
it.each([
{ models: ["grok-4.20", "grok-4.20-0309"], expected: "grok-4.20" },
{ models: ["grok-4.20-0309", "grok-4.20-0410"], expected: "grok-4.20-0410" },
{ models: ["grok-4.20-0309", "grok-4.21"], expected: "grok-4.21" },
])("orders Grok releases and short-date snapshots: $expected", ({ models, expected }) => {
const available = resolve(
deployment("everyday", "openai/gpt-6-luna"),
...models.map((model) => deployment(model, `xai/${model}`)),
);
expect(names(recommendSolverPair(available, "balanced", "grok"))).toEqual(["everyday", expected]);
});
it("does not rank a major-release snapshot above a newer minor release", () => {
const models = resolve(
deployment("old", "anthropic/claude-5-haiku-20260101"),
deployment("next", "anthropic/claude-haiku-5-1"),
deployment("hard", "anthropic/claude-opus-5"),
);
expect(names(recommendSolverPair(models, "balanced", "claude"))).toEqual(["next", "hard"]);
});
it("prefers a same-family substitute over a matching rank from another family", () => {
const models = [
...ladder().filter((model) => model.line !== "haiku"),
...resolve(deployment("luna", "openai/gpt-6-luna")),
];
const pair = recommendSolverPair(models, "balanced", "claude");
expect(names(pair)).toEqual(["medium", "large"]);
expect(pair?.notice).toContain("closest available pair");
});
it("prefers a complete family for Auto and explains a cross-family substitute when explicitly requested", () => {
const models = [
...ladder().filter((model) => model.line === "sonnet"),
...resolve(
deployment("luna", "bedrock/us.openai.gpt-6-luna"),
deployment("astra", "bedrock/us.openai.gpt-6-astra"),
),
];
expect(names(recommendSolverPair(models, "balanced", "auto"))).toEqual(["luna", "astra"]);
const selected = recommendSolverPair(models, "quality", "claude");
expect(names(selected)).toEqual(["medium", "astra"]);
expect(selected?.notice).toContain("two model families");
});
it.each([
["glm", "zai/glm-7.3-flash", "zai/glm-7.3"],
["deepseek", "deepseek/deepseek-v9-flash", "deepseek/deepseek-v9-pro"],
] as const)("handles the %s two-line family for every goal", (family, small, large) => {
const models = resolve(deployment("small", small), deployment("large", large));
for (const goal of ["cheapest", "balanced", "quality"] as const)
expect(names(recommendSolverPair(models, goal, family))).toEqual(["small", "large"]);
});
it("does not fabricate a second solver or substitute a different family when none of the requested family exists", () => {
expect(recommendSolverPair(ladder().slice(0, 1), "balanced", "auto")).toBeNull();
expect(recommendSolverPair(ladder(), "balanced", "openai")).toBeNull();
expect(recommendSolverPair([], "balanced", "auto")).toBeNull();
});
it("retains Gemini family suggestions and recognizes the existing Grok fallback", () => {
const models = resolve(
deployment("lite", "gemini/gemini-3.5-flash-lite"),
deployment("flash-old", "vertex_ai/gemini-3.5-flash"),
deployment("flash", "vertex_ai/gemini-3.8-flash"),
deployment("pro", "gemini/gemini-3.1-pro-preview"),
deployment("grok", "xai/grok-4.6"),
);
expect(names(recommendSolverPair(models, "cheapest", "gemini"))).toEqual(["lite", "flash"]);
expect(names(recommendSolverPair(models, "balanced", "gemini"))).toEqual(["lite", "pro"]);
expect(names(recommendSolverPair(models, "quality", "gemini"))).toEqual(["flash", "pro"]);
expect(models.find((m) => m.name === "grok")).toMatchObject({ family: "grok", provider: "xAI" });
expect(names(recommendSolverPair(models, "balanced", "grok"))).toEqual(["lite", "grok"]);
});
it("uses the latest deployed Terra as a fallback without replacing the preferred OpenAI pairs", () => {
const models = resolve(
deployment("terra-old", "openai/gpt-5.6-terra"),
deployment("terra", "openai/gpt-6.1-terra"),
deployment("sol", "openai/gpt-6.1-sol"),
);
expect(names(recommendSolverPair(models, "balanced", "openai"))).toEqual(["terra", "sol"]);
const complete = [
...models,
...resolve(deployment("luna", "openai/gpt-6-luna"), deployment("astra", "openai/gpt-6-astra")),
];
expect(names(recommendSolverPair(complete, "cheapest", "openai"))).toEqual(["luna", "sol"]);
expect(names(recommendSolverPair(complete, "balanced", "openai"))).toEqual(["luna", "astra"]);
expect(names(recommendSolverPair(complete, "quality", "openai"))).toEqual(["sol", "astra"]);
});
it("suggests a deterministic equal-rank cross-family fallback only when no ordered pair exists", () => {
const models = resolve(deployment("opus", "anthropic/claude-opus-5-5"), deployment("astra", "openai/gpt-6-astra"));
const pair = recommendSolverPair(models, "balanced", "auto");
expect(names(pair)).toEqual(["astra", "opus"]);
expect(pair?.notice).toContain("two model families");
expect(names(recommendSolverPair([...models].reverse(), "balanced", "auto"))).toEqual(names(pair));
const withLuna = [...models, ...resolve(deployment("luna", "openai/gpt-6-luna"))];
expect(names(recommendSolverPair(withLuna, "balanced", "auto"))).toEqual(["luna", "astra"]);
});
it("does not prefer native-provider deployments over a newer cloud deployment", () => {
const models = resolve(
deployment("native", "openai/gpt-6-luna"),
deployment("cloud", "bedrock/global.openai.gpt-6.2-luna"),
deployment("strong", "bedrock/global.openai.gpt-6-astra"),
);
expect(names(recommendSolverPair(models, "balanced", "openai"))).toEqual(["cloud", "strong"]);
});
});
describe("solverPairConfig", () => {
it("builds two solver references using built-in rules without paid classifiers, custom tiers or output overrides", () => {
expect(solverPairConfig("everyday", "hard")).toEqual({
classifier_type: "heuristic",
classification_mode: "user_turn",
tiers: { SIMPLE: ["everyday"], MEDIUM: ["everyday"], COMPLEX: ["hard"], REASONING: ["hard"] },
});
});
it("leaves the draft alone when no recommended models are available", () => {
expect(prefillEmptyTiers(empty, null)).toBe(empty);
});
describe("simple setup validation", () => {
it("avoids existing names without requiring a naming decision", () => {
expect(nextRouterName([])).toBe("auto-router");
expect(nextRouterName(["auto-router", "auto-router-2", "auto-router-4"])).toBe("auto-router-3");
});
it("preserves a partial manual assignment", () => {
const partial = { ...empty, tiers: { ...empty.tiers, SIMPLE: ["manual"] } };
expect(prefillEmptyTiers(partial, automatic)).toBe(partial);
});
it("preserves a custom tier set created before models load", () => {
const custom = { ...empty, custom_tier_set: { tiers: [], fallback_tier_id: "" } };
expect(prefillEmptyTiers(custom, automatic)).toBe(custom);
});
it.each(["capability", "llm_v2"] as const)("does not prefill %s solver pools", (classifier_type) => {
const solver = { ...empty, classifier_type };
expect(prefillEmptyTiers(solver, automatic)).toBe(solver);
it("blocks stale or incomplete inventory and same-line selections", () => {
const [small, large] = [ladder()[0], ladder()[2]];
expect(simpleSetupError(true, false, small, large)).not.toBeNull();
expect(simpleSetupError(false, true, small, large)).not.toBeNull();
expect(simpleSetupError(false, false, small)).not.toBeNull();
expect(simpleSetupError(false, false, small, small)).not.toBeNull();
expect(simpleSetupError(false, false, small, large)).toBeNull();
});
});

View file

@ -1,110 +1,263 @@
import { activeTierRows } from "./tier_rows";
import { isForecastClassifier } from "./forecast_classifier_config";
import { isAutoRouterDeployment, type AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels";
import type { AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels";
import { isAutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels";
import type { ModelGroup } from "@/components/llm_calls/fetch_models";
import { resolveAvailableModel, type AutoRouterPreset, type ModelAvailability } from "@/lib/autorouter_presets";
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
const TIER_NAMES = ["SIMPLE", "MEDIUM", "COMPLEX", "REASONING"] as const;
type TierName = (typeof TIER_NAMES)[number];
export type PreferredTierModels = Record<TierName, string[]>;
export const ROUTER_FAMILIES = {
claude: "Claude",
openai: "OpenAI",
gemini: "Gemini",
grok: "Grok",
glm: "GLM",
deepseek: "DeepSeek",
} as const;
export const ROUTER_GOALS = { cheapest: "Cheapest", balanced: "Balanced", quality: "Best quality" } as const;
export type RouterFamily = keyof typeof ROUTER_FAMILIES;
export type RouterGoal = keyof typeof ROUTER_GOALS;
export type FamilyPreference = RouterFamily | "auto";
const REASONING_EFFORT_STRENGTH = ["max", "xhigh", "high", "medium", "low", "minimal", "none"] as const;
type ModelIdentity = {
family: RouterFamily;
line: string;
rank: number;
version: readonly number[];
snapshot: number;
model: string;
};
export type SolverModel = Partial<ModelIdentity> & { name: string; provider: string; model: string };
export type RouterModel = ModelIdentity & SolverModel;
export type SolverPair = { everyday: RouterModel; hard: RouterModel; notice: string | null };
const CURRENT_TIER_MODELS: PreferredTierModels = {
SIMPLE: ["gpt-5.6-luna", "claude-haiku-4-5", "gemini-3.5-flash-lite", "deepseek-v4-flash"],
MEDIUM: ["gpt-5.6-terra", "claude-sonnet-5", "gemini-3.8-flash", "deepseek-v4-flash"],
COMPLEX: ["gpt-6-astra", "gpt-5.6-sol", "claude-opus-5", "gemini-3.1-pro-preview", "deepseek-v4-pro", "grok-4.6"],
REASONING: ["gpt-6-astra", "gpt-5.6-sol", "claude-opus-5", "gemini-3.1-pro-preview", "deepseek-v4-pro", "grok-4.6"],
const releaseVersion = (raw: string, date?: string) => {
const snapshot = raw.match(/[-.](\d{8})$/);
const release = snapshot ? raw.slice(0, -9) : raw;
return {
version: release.split(/[.-]/).map(Number),
// An undated alias follows the current snapshot of its release.
snapshot: date || snapshot ? Number((date ?? snapshot![1]).replaceAll("-", "")) : Number.MAX_SAFE_INTEGER,
};
};
export const buildPreferredTierModels = (
presets: AutoRouterPreset[],
availability: ModelAvailability,
): PreferredTierModels =>
Object.fromEntries(
TIER_NAMES.map((tier) => [
tier,
Array.from(
new Set(
[
...CURRENT_TIER_MODELS[tier],
...presets.flatMap((preset) => preset.complexity_router_config.tiers[tier]),
].flatMap((model) => {
const resolved = resolveAvailableModel(model, availability);
return resolved ? [resolved] : [];
}),
),
const parseModel = (raw: string): ModelIdentity | null => {
const model = raw
.toLowerCase()
.split("/")
.at(-1)!
.replace(/^(?:us|eu|apac|global|au|jp|us-gov)\./, "")
.replace(/^(?:anthropic|openai|zai|deepseek)\./, "")
.replace(/^databricks-/, "")
.replace(/@default$/, "")
.replace(/-v\d+(?::\d+)?$/, "");
const versionFirst = model.match(/^claude-(\d+(?:[.-]\d+)*)-(haiku|sonnet|opus)(?:[-@]([0-9]{8}))?$/);
const claude =
model.match(/^claude-(haiku|sonnet|opus)-(\d+(?:[.-]\d+)*)(?:@([0-9]+))?$/) ??
(versionFirst ? [versionFirst[0], versionFirst[2], versionFirst[1], versionFirst[3]] : null);
if (claude)
return {
family: "claude",
line: claude[1],
rank: ["haiku", "sonnet", "opus"].indexOf(claude[1]),
...releaseVersion(claude[2], claude[3]),
model,
};
const openai = model.match(/^gpt-(\d+(?:[.-]\d+)*)-(luna|terra|sol|astra)(?:-(\d{4}-\d{2}-\d{2}))?$/);
if (openai)
return {
family: "openai",
line: openai[2],
rank: openai[2] === "terra" ? 0.5 : ["luna", "sol", "astra"].indexOf(openai[2]),
...releaseVersion(openai[1], openai[3]),
model,
};
const gemini = model.match(/^gemini-(\d+(?:[.-]\d+)*)-(flash-lite|flash|pro)(?:-preview)?$/);
if (gemini)
return {
family: "gemini",
line: gemini[2],
rank: ["flash-lite", "flash", "pro"].indexOf(gemini[2]),
...releaseVersion(gemini[1]),
model,
};
const grok = model.match(/^grok-(\d+(?:[.-]\d+)*?)(?:-(\d{4}))?$/);
if (grok) return { family: "grok", line: "grok", rank: 2, ...releaseVersion(grok[1], grok[2]), model };
const glm = model.match(/^glm-(\d+(?:[.-]\d+)*)(-flash)?$/);
if (glm)
return {
family: "glm",
line: glm[2] ? "flash" : "glm",
rank: glm[2] ? 0 : 2,
...releaseVersion(glm[1]),
model,
};
const deepseek = model.match(/^deepseek-v(\d+(?:[.-]\d+)*)-(flash|pro)$/);
if (deepseek)
return {
family: "deepseek",
line: deepseek[2],
rank: deepseek[2] === "flash" ? 0 : 2,
...releaseVersion(deepseek[1]),
model,
};
return null;
};
const identityOf = (deployment: AutoRouterDeployment): ModelIdentity | null =>
parseModel(deployment.litellm_params?.model ?? "") ??
parseModel(deployment.model_info?.base_model ?? deployment.litellm_params?.base_model ?? "");
const providerOf = (deployment: AutoRouterDeployment, group: ModelGroup, family?: RouterFamily): string => {
const raw = deployment.litellm_params?.model ?? "";
const provider =
deployment.litellm_params?.custom_llm_provider ||
(raw.includes("/") ? raw.split("/")[0] : group.providers?.join(", ")) ||
(family
? { claude: "anthropic", glm: "zai", openai: "openai", gemini: "gemini", grok: "xai", deepseek: "deepseek" }[
family
]
: "");
if (provider.startsWith("bedrock")) return "Bedrock";
return (
(
{
anthropic: "Anthropic",
openai: "OpenAI",
gemini: "Google AI",
xai: "xAI",
azure: "Azure",
azure_ai: "Azure AI",
databricks: "Databricks",
github_copilot: "GitHub Copilot",
"vertex_ai-anthropic_models": "Vertex AI",
vertex_ai: "Vertex AI",
zai: "Z.ai",
deepseek: "DeepSeek",
} as Record<string, string>
)[provider] ?? provider
);
};
const compareVersion = (left: readonly number[], right: readonly number[]): number => {
const index = Array.from({ length: Math.max(left.length, right.length) }, (_, i) => i).find(
(i) => (left[i] ?? 0) !== (right[i] ?? 0),
);
return index === undefined ? 0 : (left[index] ?? 0) - (right[index] ?? 0);
};
const compareRelease = (left: ModelIdentity, right: ModelIdentity): number =>
compareVersion(left.version, right.version) || left.snapshot - right.snapshot;
export const deployedRouterModels = (
groups: readonly ModelGroup[],
deployments: readonly AutoRouterDeployment[],
): SolverModel[] =>
groups
.flatMap((group) => {
if (group.mode && group.mode !== "chat") return [];
const rows = deployments.filter((row) => row.model_name === group.model_group);
if (
rows.some(isAutoRouterDeployment) ||
group.model_group.includes("*") ||
group.model_group.startsWith("auto_router/")
)
return [];
const active = rows.filter((row) => !row.model_info?.blocked);
if (rows.length > 0 && active.length === 0) return [];
if (active.some((row) => row.model_info?.mode && row.model_info.mode !== "chat")) return [];
if (active.some((row) => row.litellm_params?.model?.includes("*"))) return [];
const identities = active.length > 0 ? active.map(identityOf) : [parseModel(group.model_group)];
const first = identities[0];
const homogeneous =
first &&
identities.every((identity) => {
if (!identity) return false;
const sameLine = identity.family === first.family && identity.line === first.line;
return sameLine && compareRelease(identity, first) === 0;
});
const identity = homogeneous ? first : undefined;
const providers =
active.length > 0
? active.map((row) => providerOf(row, group, identity?.family))
: [providerOf({}, group, identity?.family)];
return [
{
...identity,
name: group.model_group,
model: identity?.model ?? group.model_group,
provider: [...new Set(providers)].toSorted().join(", "),
},
];
})
.toSorted((a, b) => a.name.localeCompare(b.name));
const latestLines = (models: readonly RouterModel[]): RouterModel[] =>
models.filter(
(model) =>
!models.some(
(other) => other.family === model.family && other.line === model.line && compareRelease(other, model) > 0,
),
]),
) as PreferredTierModels;
const selectPreferredTierModels = (
preferredByTier: PreferredTierModels,
usableNames: ReadonlySet<string>,
): [string, string, string, string] | null => {
const preferred = TIER_NAMES.map((tier) => preferredByTier[tier].find((name) => usableNames.has(name)));
const candidates = preferred.flatMap((model, tier) => (model ? [{ model, tier }] : []));
if (candidates.length === 0) return null;
const nearest = (tier: number): string =>
[...candidates].sort(
(left, right) => Math.abs(left.tier - tier) - Math.abs(right.tier - tier) || left.tier - right.tier,
)[0].model;
return preferred.map((model, tier) => model ?? nearest(tier)) as [string, string, string, string];
};
export const buildAutomaticRouterConfig = (
models: ModelGroup[],
deployments: AutoRouterDeployment[],
preferredByTier: PreferredTierModels,
): ComplexityRouterConfigValue | null => {
const autoRouterNames: ReadonlySet<string> = new Set(
deployments
.filter(isAutoRouterDeployment)
.flatMap((deployment) => (deployment.model_name ? [deployment.model_name] : [])),
);
const names = Array.from(
new Set(
models
.filter((model) => model.mode === undefined || model.mode === "chat")
.map((model) => model.model_group)
.filter((name) => name && !name.startsWith("auto_router/") && !autoRouterNames.has(name)),
),
const requestedRanks = (goal: RouterGoal): readonly [number, number] =>
(({ cheapest: [0, 1], quality: [1, 2], balanced: [0, 2] }) as const)[goal];
export const recommendSolverPair = (
models: readonly SolverModel[],
goal: RouterGoal,
family: FamilyPreference,
): SolverPair | null => {
const recognized = models.filter((model): model is RouterModel => model.family !== undefined);
const latest = latestLines(recognized);
const [low, high] = requestedRanks(goal);
const pairs = latest.flatMap((everyday) =>
latest
.filter((hard) => {
const matchesFamily = family === "auto" || everyday.family === family || hard.family === family;
const ordered = hard.rank > everyday.rank;
const equalRankFallback = hard.rank === everyday.rank && hard.family !== everyday.family;
return (ordered || equalRankFallback) && matchesFamily;
})
.map((hard) => ({ everyday, hard })),
);
if (names.length === 0) return null;
const usableNames: ReadonlySet<string> = new Set(names);
const selected = selectPreferredTierModels(preferredByTier, usableNames);
if (selected === null) return null;
const supportedReasoningEfforts = models.find(
(model) => model.model_group === selected[3],
)?.supported_reasoning_efforts;
const reasoningEffort = REASONING_EFFORT_STRENGTH.find((effort) => supportedReasoningEfforts?.includes(effort));
return {
tiers: {
SIMPLE: [selected[0]],
MEDIUM: [selected[1]],
COMPLEX: [selected[2]],
REASONING: [selected[3]],
},
classifier_type: "heuristic_v2",
...(reasoningEffort && {
tier_model_params: { REASONING: { [selected[3]]: { reasoning_effort: reasoningEffort } } },
}),
};
const score = ({ everyday, hard }: { everyday: RouterModel; hard: RouterModel }): number =>
(everyday.family === hard.family ? 0 : 100) +
(everyday.rank === hard.rank ? 10 : 0) +
Math.abs(everyday.rank - low) +
Math.abs(hard.rank - high);
const chosen = pairs.toSorted(
(a, b) =>
score(a) - score(b) || a.everyday.name.localeCompare(b.everyday.name) || a.hard.name.localeCompare(b.hard.name),
)[0];
if (!chosen) return null;
const mixed = chosen.everyday.family !== chosen.hard.family;
const substituted = chosen.everyday.rank !== low || chosen.hard.rank !== high;
const mixedNotice = mixed ? "A complete pair is unavailable in one family. Using two model families." : null;
const substituteNotice = substituted ? "Using the closest available pair in this family." : null;
return { ...chosen, notice: mixedNotice ?? substituteNotice };
};
export const prefillEmptyTiers = (
current: ComplexityRouterConfigValue,
automatic: ComplexityRouterConfigValue | null,
): ComplexityRouterConfigValue => {
if (!automatic || current.custom_tier_set || isForecastClassifier(current.classifier_type)) return current;
if (activeTierRows(current).some((row) => row.models.length > 0)) return current;
return {
...current,
tiers: { ...current.tiers, ...automatic.tiers },
tier_model_params: { ...current.tier_model_params, ...automatic.tier_model_params },
};
export const solverPairConfig = (everyday: string, hard: string): ComplexityRouterConfigValue => ({
classifier_type: "heuristic",
classification_mode: "user_turn",
tiers: { SIMPLE: [everyday], MEDIUM: [everyday], COMPLEX: [hard], REASONING: [hard] },
});
export const simpleSetupError = (
loading: boolean,
failed: boolean,
everyday?: SolverModel,
hard?: SolverModel,
): string | null => {
if (loading) return "Loading deployed models";
if (failed) return "Could not verify deployed models";
if (!everyday || !hard) return "Select two different deployed models";
if (everyday.name === hard.name) return "Select two different deployed models";
return null;
};
export const nextRouterName = (names: readonly string[]): string => {
const candidates = [
"auto-router",
...Array.from({ length: names.length + 1 }, (_, index) => `auto-router-${index + 2}`),
];
return candidates.find((name) => !names.includes(name))!;
};

Some files were not shown because too many files have changed in this diff Show more