mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
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:
commit
447ae4d472
106 changed files with 6013 additions and 2817 deletions
|
|
@ -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()
|
||||
|
|
@ -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 }}
|
||||
65
.github/workflows/sync-together-ai-models.yml
vendored
65
.github/workflows/sync-together-ai-models.yml
vendored
|
|
@ -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 }}
|
||||
|
|
@ -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.
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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__"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 ""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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": {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 }
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -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 });
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
638
tests/unit/llms/chatgpt/chat/test_chatgpt_chat_transformation.py
Normal file
638
tests/unit/llms/chatgpt/chat/test_chatgpt_chat_transformation.py
Normal 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
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
90
tests/unit/llms/chatgpt/test_chatgpt_common_utils.py
Normal file
90
tests/unit/llms/chatgpt/test_chatgpt_common_utils.py
Normal 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}")
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -1086,6 +1086,7 @@ _COST_MAP_ROW_SCHEMA: Final = {
|
|||
"chat",
|
||||
"completion",
|
||||
"container",
|
||||
"decisions",
|
||||
"image_edit",
|
||||
"embedding",
|
||||
"evaluation",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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} />);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 () => {
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
);
|
||||
}
|
||||
|
|
@ -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>
|
||||
);
|
||||
}
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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
Loading…
Add table
Reference in a new issue