chore: merge origin/main into test cleanup

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-16 00:33:25 +00:00
commit ba6bcd747f
145 changed files with 9064 additions and 507 deletions

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.67"
version = "0.1.68"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.67"
version = "0.1.68"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

View file

@ -96,6 +96,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/langfuse/",
"/vllm/",
"/mistral/",
"/nvidia_nim/",
"/groq/",
"/voyage/",
"/cursor/",

View file

@ -40,4 +40,4 @@ if not logger.handlers:
logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
logger.setLevel(os.getenv("LITELLM_LOG", "INFO").upper())

View file

@ -0,0 +1,18 @@
-- DropIndex
DROP INDEX IF EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_is_act_idx";
-- DropIndex
DROP INDEX IF EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_key";
-- AlterTable
-- NOT NULL DEFAULT '' (not nullable): Postgres unique constraints treat every
-- NULL as distinct, so a nullable column would let multiple unscoped mappings
-- collide on the same claim without a constraint violation. The constant
-- default is a fast, metadata-only backfill for existing rows, not a rewrite.
ALTER TABLE "LiteLLM_JWTKeyMapping" ADD COLUMN IF NOT EXISTS "jwt_issuer" TEXT NOT NULL DEFAULT '';
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_issuer_jwt_claim_name_jwt_claim_v_idx" ON "LiteLLM_JWTKeyMapping"("jwt_issuer", "jwt_claim_name", "jwt_claim_value", "is_active");
-- CreateIndex
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_issuer_jwt_claim_name_jwt_claim_v_key" ON "LiteLLM_JWTKeyMapping"("jwt_issuer", "jwt_claim_name", "jwt_claim_value");

View file

@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
model LiteLLM_JWTKeyMapping {
id String @id @default(uuid())
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
// Not nullable: Postgres unique constraints treat every NULL as
// distinct, so a nullable column would let multiple unscoped
// mappings collide on the same claim without a constraint violation.
jwt_claim_name String // e.g. "sub", "email"
jwt_claim_value String // The claim value to match
token String // Hashed virtual key (FK)
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
}
// Deprecated keys during grace period - allows old key to work until revoke_at

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.97"
version = "0.4.98"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.97"
version = "0.4.98"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

@ -343,6 +343,7 @@ _anthropic_prompt_caching_ttl_env: Optional[str] = os.getenv("LITELLM_ANTHROPIC_
anthropic_prompt_caching_ttl: Optional[Literal["5m", "1h"]] = (
"1h" if _anthropic_prompt_caching_ttl_env == "1h" else "5m" if _anthropic_prompt_caching_ttl_env == "5m" else None
)
openai_system_messages_first: bool = False
disable_vertex_batch_output_transformation: bool = False
extra_spend_tag_headers: Optional[List[str]] = None
in_memory_llm_clients_cache: "LLMClientCache"

View file

@ -401,10 +401,14 @@ def _parse_json_logs_env(value: str | None) -> bool:
return (value or "").lower() == "true"
def resolve_log_level(log_level: str) -> int:
return getattr(logging, log_level.upper())
json_logs: Final = _parse_json_logs_env(os.getenv("JSON_LOGS"))
# Create a handler for the logger (you may need to adapt this based on your needs)
log_level: Final = os.getenv("LITELLM_LOG", "DEBUG")
numeric_level: Final[str] = getattr(logging, log_level.upper())
numeric_level: Final[int] = resolve_log_level(log_level)
handler: Final = LevelRoutingStreamHandler()
handler.setLevel(numeric_level)
handler.addFilter(_secret_filter)

View file

@ -317,6 +317,8 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update"
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure"
CLIENT_REQUESTED_MODEL_SCOPE_KEY: Final = "litellm.client_requested_model"
MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY: Final = "litellm.model_group_alias_resolved"
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged"
@ -1776,6 +1778,7 @@ DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_PROMPT_
LENGTH_OF_LITELLM_GENERATED_KEY: Final = int(os.getenv("LENGTH_OF_LITELLM_GENERATED_KEY", 16))
MINIMUM_CUSTOM_KEY_LENGTH: Final = int(os.getenv("MINIMUM_CUSTOM_KEY_LENGTH", 16))
SECRET_MANAGER_REFRESH_INTERVAL: Final = int(os.getenv("SECRET_MANAGER_REFRESH_INTERVAL", 86400))
OPENAI_SYSTEM_MESSAGES_FIRST_PROVIDERS: Final = frozenset({"openai", "azure"})
LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
"default_internal_user_params",
"default_team_params",
@ -1793,6 +1796,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
# test_general_settings_ui_fields_are_db_overridable enforces that pairing.
"enable_anthropic_prompt_caching",
"anthropic_prompt_caching_ttl",
"openai_system_messages_first",
"max_ui_session_budget",
"budget_rollover",
"mcp_tool_search",

View file

@ -1203,6 +1203,24 @@ def _without_provider_stated_cost(usage: Usage | None) -> Usage | None:
return usage.model_copy(update=MappingProxyType({"cost": None}))
def _split_responses_ws_logging_object_by_service_tier(
completion_response: LiteLLMRealtimeStreamLoggingObject,
) -> tuple[LiteLLMRealtimeStreamLoggingObject, ...] | None:
partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(
cast(Sequence[Mapping[str, object]], completion_response.results)
)
if len(partition) <= 1:
return None
return tuple(
LiteLLMRealtimeStreamLoggingObject(
results=cast(OpenAIRealtimeStreamList, list(group)),
usage=ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(group),
service_tier=tier,
)
for tier, group in partition.items()
)
def completion_cost(
completion_response: object | None = None,
model: str | None = None,
@ -1266,6 +1284,41 @@ def completion_cost(
try:
call_type = _infer_call_type(call_type, completion_response) or "completion"
if call_type == CallTypes.aresponses_websocket.value and isinstance(
completion_response, LiteLLMRealtimeStreamLoggingObject
):
ws_tier_parts: Final = _split_responses_ws_logging_object_by_service_tier(completion_response)
if ws_tier_parts is not None:
return sum(
completion_cost(
completion_response=part,
model=model,
prompt=prompt,
messages=messages,
completion=completion,
total_time=total_time,
call_type=call_type,
custom_llm_provider=custom_llm_provider,
region_name=region_name,
size=size,
quality=quality,
n=n,
custom_cost_per_token=custom_cost_per_token,
custom_cost_per_second=custom_cost_per_second,
optional_params=optional_params,
custom_pricing=custom_pricing,
base_model=base_model,
standard_built_in_tools_params=standard_built_in_tools_params,
litellm_model_name=litellm_model_name,
router_model_id=router_model_id,
litellm_logging_obj=litellm_logging_obj,
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
)
for part in ws_tier_parts
)
if (
(call_type == "aimage_generation" or call_type == "image_generation")
and model is not None
@ -1466,12 +1519,15 @@ def completion_cost(
duration_seconds = usage_obj.get("duration_seconds", None)
_vr = usage_obj.get("video_resolution", None)
provider_reported_cost = usage_obj.get("provider_reported_cost_usd", None)
_vc = usage_obj.get("video_count", None)
else:
duration_seconds = getattr(usage_obj, "duration_seconds", None)
_vr = getattr(usage_obj, "video_resolution", None)
provider_reported_cost = getattr(usage_obj, "provider_reported_cost_usd", None)
_vc = getattr(usage_obj, "video_count", None)
if _vr is not None:
video_resolution = str(_vr).strip().lower()
video_count = _vc if isinstance(_vc, int) and not isinstance(_vc, bool) and _vc > 1 else 1
if _video_model_info is None and provider_reported_cost is not None:
return float(provider_reported_cost)
@ -1482,12 +1538,15 @@ def completion_cost(
video_generation_cost,
)
return video_generation_cost(
model=model,
duration_seconds=duration_seconds,
custom_llm_provider=custom_llm_provider,
model_info=_video_model_info,
video_resolution=video_resolution,
return (
video_generation_cost(
model=model,
duration_seconds=duration_seconds,
custom_llm_provider=custom_llm_provider,
model_info=_video_model_info,
video_resolution=video_resolution,
)
* video_count
)
# Fallback to default video cost calculation if no duration available
return default_video_cost_calculator(
@ -2558,6 +2617,7 @@ _RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "re
class _ResponsesWsEventResponse(BaseModel):
usage: Mapping[str, object] | None = None
service_tier: str | None = None
class _ResponsesWsEvent(BaseModel):
@ -2565,20 +2625,39 @@ class _ResponsesWsEvent(BaseModel):
response: _ResponsesWsEventResponse | None = None
def _billable_responses_ws_events(
results: Sequence[Mapping[str, object]],
) -> tuple[tuple[Mapping[str, object], _ResponsesWsEventResponse], ...]:
return tuple(
(result, event.response)
for result in results
if (event := _ResponsesWsEvent.model_validate(result)).type in _RESPONSES_WS_BILLABLE_EVENT_TYPES
and event.response is not None
and event.response.usage is not None
)
class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor):
@staticmethod
def collect_usage_from_responses_ws_results(
results: Sequence[Mapping[str, object]],
) -> tuple[Usage, ...]:
events: Final = tuple(_ResponsesWsEvent.model_validate(result) for result in results)
return tuple(
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # same shared transform the realtime processor uses
event.response.usage
response.usage
)
for event in events
if event.type in _RESPONSES_WS_BILLABLE_EVENT_TYPES
and event.response is not None
and event.response.usage is not None
for _, response in _billable_responses_ws_events(results)
if response.usage is not None
)
@staticmethod
def partition_results_by_service_tier(
results: Sequence[Mapping[str, object]],
) -> Mapping[str | None, tuple[Mapping[str, object], ...]]:
billable: Final = _billable_responses_ws_events(results)
tiers: Final = dict.fromkeys(response.service_tier for _, response in billable)
return MappingProxyType(
{tier: tuple(result for result, response in billable if response.service_tier == tier) for tier in tiers}
)
@staticmethod

View file

@ -2101,9 +2101,14 @@ class Logging(LiteLLMLoggingBaseClass):
results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
)
)
ws_tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(
results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
)
ws_service_tier: Final = next(iter(ws_tier_partition)) if len(ws_tier_partition) == 1 else None
logging_result = LiteLLMRealtimeStreamLoggingObject(
usage=combined_ws_usage,
results=result, # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
service_tier=ws_service_tier,
)
elif (

View file

@ -2256,6 +2256,22 @@ def drop_tool_reference_parts_from_tool_messages(
return [_drop_tool_reference_parts(message) for message in messages] # mutable-ok: pipelines mutate message lists
INSTRUCTION_MESSAGE_ROLES: Final = frozenset({"system", "developer"})
def _is_instruction_message(message: AllMessageValues) -> bool:
return message.get("role") in INSTRUCTION_MESSAGE_ROLES
def system_messages_first(
messages: list[AllMessageValues], # mutable-ok: message pipelines type messages as mutable lists
) -> list[AllMessageValues]: # mutable-ok: message pipelines type messages as mutable lists
return [ # mutable-ok: pipelines mutate message lists
*(message for message in messages if _is_instruction_message(message)),
*(message for message in messages if not _is_instruction_message(message)),
]
def _attempt_json_repair(s: str) -> object | None:
"""
Attempt to repair truncated JSON produced by LLM tool calls.

View file

@ -1500,19 +1500,33 @@ def convert_to_gemini_tool_call_result(
return _part
def _sanitize_anthropic_tool_use_id(tool_use_id: str) -> str:
"""
Sanitize tool_use_id to match Anthropic's required pattern: ^[a-zA-Z0-9_-]+$
_TOOL_USE_ID_FALLBACK: Final = "tool_use_id"
_ANTHROPIC_TOOL_USE_ID_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]")
_BEDROCK_TOOL_USE_ID_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_.:-]")
_BEDROCK_TOOL_USE_ID_MAX_LEN: Final = 64
_BEDROCK_TOOL_USE_ID_HASH_LEN: Final = 8
Anthropic requires tool_use_id to only contain alphanumeric characters, underscores, and hyphens.
This function replaces any invalid characters with underscores.
def _replace_invalid_tool_use_id_chars(tool_use_id: str, invalid_chars: re.Pattern[str]) -> str:
return invalid_chars.sub("_", tool_use_id) or _TOOL_USE_ID_FALLBACK
def _sanitize_anthropic_tool_use_id(tool_use_id: str) -> str:
"""Anthropic requires tool_use_id to match ^[a-zA-Z0-9_-]+$."""
return _replace_invalid_tool_use_id_chars(tool_use_id, _ANTHROPIC_TOOL_USE_ID_INVALID_CHARS)
def _sanitize_bedrock_tool_use_id(tool_use_id: str) -> str:
"""
# Replace any character that's not alphanumeric, underscore, or hyphen with underscore
sanitized = re.sub(r"[^a-zA-Z0-9_-]", "_", tool_use_id)
# Ensure it's not empty (fallback to a default if needed)
if not sanitized:
sanitized = "tool_use_id"
return sanitized
Bedrock Converse requires toolUseId to match [a-zA-Z0-9_.:-]+ and be at most 64 chars.
Ids that need rewriting get a short hash of the original appended so two ids that only
differ in a replaced char or past the cut still map to distinct values.
"""
sanitized: Final = _replace_invalid_tool_use_id_chars(tool_use_id, _BEDROCK_TOOL_USE_ID_INVALID_CHARS)
if sanitized == tool_use_id and len(sanitized) <= _BEDROCK_TOOL_USE_ID_MAX_LEN:
return sanitized
digest: Final = hashlib.sha256(tool_use_id.encode()).hexdigest()[:_BEDROCK_TOOL_USE_ID_HASH_LEN]
return f"{sanitized[: _BEDROCK_TOOL_USE_ID_MAX_LEN - _BEDROCK_TOOL_USE_ID_HASH_LEN - 1]}_{digest}"
_ANTHROPIC_DOCUMENT_BASE64_MEDIA_TYPES: Final = {"application/pdf", "text/plain"}
@ -3661,7 +3675,9 @@ def _convert_to_bedrock_tool_call_invoke(
if parsed_objects:
# First object keeps the original tool id.
for obj_idx, obj in enumerate(parsed_objects):
block_id = tool_id if obj_idx == 0 else f"{tool_id}_{obj_idx}"
block_id = _sanitize_bedrock_tool_use_id(
tool_id if obj_idx == 0 else f"{tool_id}_{obj_idx}"
)
bedrock_tool = BedrockToolUseBlock(input=obj, name=name, toolUseId=block_id)
_parts_list.append(BedrockContentBlock(toolUse=bedrock_tool))
# cache_control applies to the whole original
@ -3678,7 +3694,9 @@ def _convert_to_bedrock_tool_call_invoke(
# Fallback: no objects extracted — use empty dict.
arguments_dict = {}
bedrock_tool = BedrockToolUseBlock(input=arguments_dict, name=name, toolUseId=tool_id)
bedrock_tool = BedrockToolUseBlock(
input=arguments_dict, name=name, toolUseId=_sanitize_bedrock_tool_use_id(tool_id)
)
bedrock_content_block = BedrockContentBlock(toolUse=bedrock_tool)
_parts_list.append(bedrock_content_block)
@ -3849,7 +3867,7 @@ def _convert_to_bedrock_tool_call_result(
tool_result_content_blocks, used_search_results = _build_bedrock_tool_result_content_blocks(message)
message.get("name", "")
id: Final = str(message.get("tool_call_id", str(uuid.uuid4())))
id: Final = _sanitize_bedrock_tool_use_id(str(message.get("tool_call_id", str(uuid.uuid4()))))
tool_result: Final = BedrockToolResultBlock(content=tool_result_content_blocks, toolUseId=id)
if used_search_results:

View file

@ -539,6 +539,13 @@ class AnthropicModelInfo(BaseLLMModelInfo):
value: Final = litellm.model_cost.get(model, {}).get(key)
return value if isinstance(value, bool) else None
@staticmethod
def supports_fast_mode(model: str, custom_llm_provider: str) -> bool:
return (
custom_llm_provider == "anthropic"
and AnthropicModelInfo._get_exact_model_capability(model, "supports_fast_mode") is True
)
@staticmethod
def _get_provider_resolved_capability(model: str, key: str, custom_llm_provider: str) -> bool | None:
"""Resolve boolean capability ``key`` for ``model`` under the caller's provider.

View file

@ -9,6 +9,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
drop_tool_reference_parts_from_tool_messages,
flatten_combinators_and_drop_non_python_regex_patterns,
hoist_images_from_tool_messages,
system_messages_first,
tool_with_sanitized_parameters,
)
from litellm.litellm_core_utils.prompt_templates.factory import (
@ -276,7 +277,8 @@ class AzureOpenAIConfig(BaseConfig):
litellm_params: dict,
headers: dict,
) -> dict:
stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(messages)
ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages
stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages)
azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages))
return {
"model": model,

View file

@ -5,7 +5,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Final
import httpx
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.llms.azure_ai.common_utils import (
@ -18,6 +18,8 @@ from litellm.llms.base_llm.passthrough.transformation import (
BasePassthroughConfig,
RelayShape,
logged_relay_shape,
model_group_from,
relayed_body,
strip_leading_model_segment,
)
from litellm.types.llms.openai import AllMessageValues
@ -35,19 +37,6 @@ if TYPE_CHECKING:
EMPTY_QUERY: Final[Mapping[str, object]] = MappingProxyType({})
class PassthroughMetadata(BaseModel):
model_config = ConfigDict(extra="ignore")
model_group: str = ""
def model_group_from(litellm_params: Mapping[str, object]) -> str:
try:
return PassthroughMetadata.model_validate(litellm_params.get("litellm_metadata")).model_group
except ValidationError:
return ""
def api_version_from(litellm_params: Mapping[str, object]) -> str | None:
try:
return TypeAdapter(str | None).validate_python(litellm_params.get("api_version"))
@ -96,14 +85,6 @@ def relay_query_params(
return MappingProxyType({**(request_query_params or EMPTY_QUERY), "api-version": api_version})
def relayed_body(httpx_response: Response) -> str | dict:
try:
body: Final[object] = httpx_response.json()
except ValueError:
return httpx_response.text
return body if isinstance(body, dict) else httpx_response.text
FOUNDRY_RELAY_SHAPES: Final = (
RelayShape("/rerank", CallTypes.arerank, RerankResponse.model_validate),
RelayShape("/providers/blackforestlabs/v1/flux-2-pro", CallTypes.aimage_generation, ImageResponse.model_validate),

View file

@ -6,7 +6,7 @@ from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
from pydantic import TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from litellm.types.utils import CallTypes
@ -29,6 +29,19 @@ if TYPE_CHECKING:
RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
class PassthroughMetadata(BaseModel):
model_config = ConfigDict(extra="ignore")
model_group: str = ""
def model_group_from(litellm_params: Mapping[str, object]) -> str:
try:
return PassthroughMetadata.model_validate(litellm_params.get("litellm_metadata")).model_group
except ValidationError:
return ""
def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str:
path: Final = endpoint.lstrip("/")
for model_name in model_names:
@ -55,6 +68,14 @@ def relayed_json_object(httpx_response: Response) -> Mapping[str, object] | None
return None
def relayed_body(httpx_response: Response) -> str | dict:
try:
body: Final[object] = httpx_response.json()
except ValueError:
return httpx_response.text
return body if isinstance(body, dict) else httpx_response.text
@dataclass(frozen=True, slots=True)
class RelayShape:
path_suffix: str

View file

@ -9,6 +9,7 @@ import litellm
from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS
from litellm.images.utils import ImageEditRequestUtils
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
from litellm.llms.vertex_ai.videos.transformation import veo_video_count_from_parameters
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.gemini import (
GeminiLongRunningOperationResponse,
@ -354,6 +355,9 @@ class GeminiVideoConfig(BaseVideoConfig):
video_resolution: Final = _usage_video_resolution_from_parameters(parameters)
if video_resolution is not None:
usage_data["video_resolution"] = video_resolution
video_count: Final = veo_video_count_from_parameters(parameters)
if video_count is not None:
usage_data["video_count"] = video_count
video_obj.usage = usage_data
return video_obj

View file

@ -0,0 +1,139 @@
from __future__ import annotations
import re
from collections.abc import Collection, Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Final
import httpx
from litellm.llms.base_llm.passthrough.transformation import (
BasePassthroughConfig,
model_group_from,
relayed_body,
strip_leading_model_segment,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import DeploymentTypedDict
from litellm.types.utils import LlmProviders, StandardPassThroughResponseObject
if TYPE_CHECKING:
from httpx import URL, Response
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse
API_VERSION_SEGMENT: Final = re.compile(r"^v\d+$")
NVIDIA_NIM_MODEL_PREFIX: Final = f"{LlmProviders.NVIDIA_NIM.value}/"
NVIDIA_NIM_ROUTE_PREFIX: Final = re.compile(rf"^/{LlmProviders.NVIDIA_NIM.value}/", re.IGNORECASE)
def is_nvidia_nim_deployment(deployment: DeploymentTypedDict) -> bool:
litellm_params: Final = deployment["litellm_params"]
return litellm_params.get("custom_llm_provider") == LlmProviders.NVIDIA_NIM.value or litellm_params.get(
"model", ""
).startswith(NVIDIA_NIM_MODEL_PREFIX)
def nvidia_nim_model_groups(deployments: Iterable[DeploymentTypedDict] | None) -> frozenset[str]:
listed: Final = tuple(deployments or ())
nim_groups: Final = frozenset(d["model_name"] for d in listed if is_nvidia_nim_deployment(d))
other_groups: Final = frozenset(d["model_name"] for d in listed if not is_nvidia_nim_deployment(d))
return nim_groups - other_groups
def nvidia_nim_model_group_in_path(path: str, deployments: Iterable[DeploymentTypedDict] | None) -> str | None:
return nvidia_nim_router_model_in_endpoint(
NVIDIA_NIM_ROUTE_PREFIX.sub("", path), nvidia_nim_model_groups(deployments)
)
def nvidia_nim_router_model_in_endpoint(endpoint: str, router_models: Collection[str]) -> str | None:
segments: Final = tuple(segment for segment in endpoint.split("/") if segment)
return next(
(
"/".join(segments[:length])
for length in range(len(segments), 0, -1)
if "/".join(segments[:length]) in router_models
),
None,
)
def without_repeated_version_prefix(api_base: str, native_endpoint: str) -> str:
url: Final = httpx.URL(api_base)
base_segments: Final = tuple(segment for segment in url.path.split("/") if segment)
first_native_segment: Final = native_endpoint.lstrip("/").split("/", 1)[0]
repeated: Final = (
bool(base_segments)
and API_VERSION_SEGMENT.match(first_native_segment) is not None
and base_segments[-1] == first_native_segment
)
kept_segments: Final = base_segments[:-1] if repeated else base_segments
return str(url.copy_with(path="/" + "/".join(kept_segments), query=None)).rstrip("/")
class NvidiaNimPassthroughConfig(BasePassthroughConfig):
def is_streaming_request(self, endpoint: str, request_data: dict) -> bool:
return bool(request_data.get("stream", False))
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
endpoint: str,
request_query_params: dict | None,
litellm_params: dict,
) -> tuple[URL, str]:
base_target_url: Final = self.get_api_base(api_base)
if base_target_url is None:
raise ValueError("NVIDIA NIM api base not found: set `api_base` on the deployment or NVIDIA_NIM_API_BASE")
native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params)))
root: Final = without_repeated_version_prefix(base_target_url, native_endpoint)
return (self.format_url(native_endpoint, root, request_query_params), root)
def validate_environment(
self,
headers: Mapping[str, str],
model: str,
messages: Sequence[AllMessageValues],
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
) -> dict[str, str]: # mutable-ok: base class contract returns dict for httpx
if api_key is None:
return dict(headers) # mutable-ok: base class contract returns dict for httpx
return {
**headers,
"Authorization": f"Bearer {api_key}",
} # mutable-ok: base class contract returns dict for httpx
@staticmethod
def get_api_base(api_base: str | None = None) -> str | None:
return api_base or get_secret_str("NVIDIA_NIM_API_BASE")
@staticmethod
def get_api_key(api_key: str | None = None) -> str | None:
return api_key or get_secret_str("NVIDIA_NIM_API_KEY")
@staticmethod
def get_base_model(model: str) -> str | None:
return model
def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]:
return []
def logging_non_streaming_response(
self,
model: str,
custom_llm_provider: str,
httpx_response: Response,
request_data: Mapping[str, object],
logging_obj: Logging,
endpoint: str,
) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None:
return StandardPassThroughResponseObject(response=relayed_body(httpx_response))

View file

@ -12,6 +12,7 @@ from urllib.parse import urlparse
import httpx
import litellm
from litellm.constants import OPENAI_SYSTEM_MESSAGES_FIRST_PROVIDERS
from litellm.litellm_core_utils.core_helpers import map_finish_reason
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
_extract_reasoning_content,
@ -24,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
flatten_combinators_and_drop_non_python_regex_patterns,
get_tool_call_names,
hoist_images_from_tool_messages,
system_messages_first,
tool_with_sanitized_parameters,
)
from litellm.litellm_core_utils.prompt_templates.image_handling import (
@ -463,6 +465,15 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
]
return MappingProxyType({"tools": sanitized})
def _prompt_cache_ordered_messages(
self, messages: list[AllMessageValues], litellm_params: Mapping[str, object]
) -> list[AllMessageValues]:
if not litellm.openai_system_messages_first:
return messages
if litellm_params.get("custom_llm_provider") not in OPENAI_SYSTEM_MESSAGES_FIRST_PROVIDERS:
return messages
return system_messages_first(messages)
def transform_request(
self,
model: str,
@ -477,7 +488,9 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
Returns:
dict: The transformed request. Sent as the body of the API call.
"""
messages = self._transform_messages(messages=messages, model=model)
messages = self._transform_messages(
messages=self._prompt_cache_ordered_messages(messages, litellm_params), model=model
)
if not self._should_preserve_cache_control_for_endpoint(
litellm_params.get("custom_llm_provider"), litellm_params.get("api_base")
):
@ -506,7 +519,9 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
litellm_params: dict,
headers: dict,
) -> dict:
transformed_messages = await self._transform_messages(messages=messages, model=model, is_async=True)
transformed_messages = await self._transform_messages(
messages=self._prompt_cache_ordered_messages(messages, litellm_params), model=model, is_async=True
)
if not self._should_preserve_cache_control_for_endpoint(
litellm_params.get("custom_llm_provider"), litellm_params.get("api_base")
):

View file

@ -70,10 +70,17 @@ def _parse_veo_operation(raw_response: httpx.Response) -> _VeoOperation:
return operation
def veo_video_count_from_parameters(parameters: Mapping[str, object]) -> int | None:
sample_count: Final = parameters.get("sampleCount")
if isinstance(sample_count, bool) or not isinstance(sample_count, int) or sample_count < 1:
return None
return sample_count
def _build_vertex_video_usage_from_request_data(
request_data: dict[str, Any] | None,
) -> dict[str, float | str]:
"""Build usage metadata (duration, resolution) for video cost calculation."""
"""Build usage metadata (duration, resolution, video count) for video cost calculation."""
usage_data: Final[dict[str, float | str]] = {}
if not request_data:
return usage_data
@ -88,6 +95,9 @@ def _build_vertex_video_usage_from_request_data(
res: Final = parameters.get("resolution")
if res is not None and str(res).strip() != "":
usage_data["video_resolution"] = str(res).strip().lower()
video_count: Final = veo_video_count_from_parameters(parameters)
if video_count is not None:
usage_data["video_count"] = video_count
return usage_data

File diff suppressed because it is too large Load diff

View file

@ -205,6 +205,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
"/gigachat/",
"/milvus/",
"/mistral/",
"/nvidia_nim/",
"/openai/",
"/openai_passthrough/",
"/vertex-ai/",

View file

@ -15226,6 +15226,17 @@
"title": "Jwt Claim Value",
"type": "string"
},
"jwt_issuer": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Jwt Issuer"
},
"key": {
"title": "Key",
"type": "string"
@ -15310,6 +15321,17 @@
"title": "Jwt Claim Value",
"type": "string"
},
"jwt_issuer": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Jwt Issuer"
},
"updated_at": {
"format": "date-time",
"title": "Updated At",
@ -15366,6 +15388,17 @@
],
"title": "Is Active"
},
"jwt_issuer": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Jwt Issuer"
},
"key": {
"anyOf": [
{
@ -18912,6 +18945,228 @@
]
}
},
"/nvidia_nim/{endpoint}": {
"delete": {
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__delete",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"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": "Nvidia Nim Proxy Route",
"tags": [
"llm_passthrough"
]
},
"get": {
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__get",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"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": "Nvidia Nim Proxy Route",
"tags": [
"llm_passthrough"
]
},
"patch": {
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__patch",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"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": "Nvidia Nim Proxy Route",
"tags": [
"llm_passthrough"
]
},
"post": {
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__post",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"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": "Nvidia Nim Proxy Route",
"tags": [
"llm_passthrough"
]
},
"put": {
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__put",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"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": "Nvidia Nim Proxy Route",
"tags": [
"llm_passthrough"
]
}
},
"/openai/deployments/{model}/chat/completions": {
"post": {
"description": "Follows the exact same API spec as `OpenAI's Chat API https://platform.openai.com/docs/api-reference/chat`\n\n```bash\ncurl -X POST http://localhost:4000/v1/chat/completions \n-H \"Content-Type: application/json\" \n-H \"Authorization: Bearer sk-1234\" \n-d '{\n \"model\": \"gpt-4o\",\n \"messages\": [\n {\n \"role\": \"user\",\n \"content\": \"Hello!\"\n }\n ]\n}'\n```",

View file

@ -246,6 +246,7 @@ class Litellm_EntityType(enum.Enum):
TEAM = "team"
TEAM_MEMBER = "team_member"
ORGANIZATION = "organization"
ORGANIZATION_MEMBER = "organization_member"
PROJECT = "project"
TAG = "tag"
AGENT = "agent"
@ -485,6 +486,7 @@ class LiteLLMRoutes(enum.Enum):
"/milvus",
"/gigachat",
"/watsonx",
"/nvidia_nim",
]
#########################################################
@ -4486,12 +4488,14 @@ class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
jwt_claim_name: str
jwt_claim_value: str
key: str
jwt_issuer: str | None = None
description: str | None = None
class UpdateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
id: str
key: str | None = None
jwt_issuer: str | None = None
description: str | None = None
is_active: bool | None = None
@ -4502,6 +4506,7 @@ class DeleteJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
class JWTKeyMappingResponse(LiteLLMPydanticObjectBase):
id: str
jwt_issuer: str | None = None
jwt_claim_name: str
jwt_claim_value: str
description: str | None = None
@ -5253,6 +5258,7 @@ class DBSpendUpdateTransactions(TypedDict):
team_list_transactions: dict[str, float] | None
team_member_list_transactions: dict[str, float] | None
org_list_transactions: dict[str, float] | None
org_member_list_transactions: ReadOnly[dict[str, float] | None]
tag_list_transactions: dict[str, float] | None
agent_list_transactions: dict[str, float] | None
model_access_group_list_transactions: ReadOnly[dict[str, float] | None]

View file

@ -148,6 +148,7 @@ class _PrismaDictableRow(Protocol):
class _PrismaJWTKeyMappingRow(Protocol):
token: str
jwt_issuer: str
jwt_claim_name: str
jwt_claim_value: str
@ -3601,9 +3602,18 @@ async def _fetch_key_object_from_db_with_reconnect(
raise
def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str) -> str:
"""Cache key under which ``_resolve_jwt_to_virtual_key`` stores a JWT-claim-to-key mapping."""
return f"jwt_key_mapping:{jwt_claim_name}:{jwt_claim_value}"
def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str, jwt_issuer: str | None = None) -> str:
"""Cache key under which a JWT-claim-to-key mapping is stored, scoped to one
issuer (or the issuer-agnostic/global scope when ``jwt_issuer`` is falsy).
Scoped by issuer (when one is configured) so a cached hit or ``__NO_MAPPING__`` miss
for one issuer's claim value can never be served to a different issuer whose claim
value happens to collide. Unchanged for the global scope, keeping the single-issuer
(no ``litellm_jwtauth.issuers`` configured) cache key format stable across this fix.
"""
if not jwt_issuer:
return f"jwt_key_mapping:{jwt_claim_name}:{jwt_claim_value}"
return f"jwt_key_mapping:{jwt_issuer}:{jwt_claim_name}:{jwt_claim_value}"
@log_db_metrics
@ -3615,7 +3625,7 @@ async def get_jwt_key_mapping_cache_keys_for_token(
mappings: Final = await _jwt_key_mapping_table(JWTKeyMappingRepository(prisma_client)).find_many(
where={"token": hashed_token}
)
return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value) for m in mappings)
return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value, m.jwt_issuer) for m in mappings)
@log_db_metrics
@ -3623,9 +3633,14 @@ async def get_jwt_key_mapping_object(
jwt_claim_name: str,
jwt_claim_value: str,
prisma_client: PrismaClient,
jwt_issuer: str | None = None,
) -> str | None:
"""
Lookup a JWT-to-virtual-key mapping from the database.
Lookup a JWT-to-virtual-key mapping from the database for one exact scope:
``jwt_issuer`` (or the global/issuer-agnostic scope when falsy). Does not fall
back to the global scope itself -- a caller that wants "issuer-scoped mapping,
else the global one" queries both scopes itself, so each result can be cached
under its own scope's key (see ``_resolve_jwt_to_virtual_key``).
Returns the hashed token (str) if a matching active mapping is found, else None.
"""
@ -3633,6 +3648,7 @@ async def get_jwt_key_mapping_object(
where={
"jwt_claim_name": jwt_claim_name,
"jwt_claim_value": jwt_claim_value,
"jwt_issuer": jwt_issuer or "",
"is_active": True,
}
)

View file

@ -28,6 +28,7 @@ from litellm.litellm_core_utils.url_utils import (
validate_url,
)
from litellm.llms.azure.passthrough.transformation import azure_router_model_in_endpoint
from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path
from litellm.proxy._types import *
from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_metadata
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
@ -976,6 +977,26 @@ def _get_deployment_default_tpm_limit(model_name: str) -> int | None:
return _get_deployment_default_limit(model_name, "default_api_key_tpm_limit")
def get_key_own_model_rate_limit(
user_api_key_dict: UserAPIKeyAuth,
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
) -> dict[str, int] | None:
if user_api_key_dict.metadata:
result: Final = user_api_key_dict.metadata.get(rate_limit_key)
if result:
return result
if not user_api_key_dict.model_max_budget:
return None
budget_key: Final = "rpm_limit" if rate_limit_key == "model_rpm_limit" else "tpm_limit"
model_limit: Final = {
model: budget[budget_key]
for model, budget in user_api_key_dict.model_max_budget.items()
if isinstance(budget, dict) and budget.get(budget_key) is not None
}
return model_limit or None
def get_key_model_rpm_limit(
user_api_key_dict: UserAPIKeyAuth,
model_name: str | None = None,
@ -989,20 +1010,9 @@ def get_key_model_rpm_limit(
3. Team metadata (model_rpm_limit)
4. Deployment default_api_key_rpm_limit (when model_name is provided)
"""
# 1. Check key metadata first (takes priority)
if user_api_key_dict.metadata:
result: Final = user_api_key_dict.metadata.get("model_rpm_limit")
if result:
return result
# 2. Check model_max_budget
if user_api_key_dict.model_max_budget:
model_rpm_limit: Final[dict[str, int]] = {}
for model, budget in user_api_key_dict.model_max_budget.items():
if isinstance(budget, dict) and budget.get("rpm_limit") is not None:
model_rpm_limit[model] = budget["rpm_limit"]
if model_rpm_limit:
return model_rpm_limit
key_own_limit: Final = get_key_own_model_rate_limit(user_api_key_dict, "model_rpm_limit")
if key_own_limit is not None:
return key_own_limit
# 3. Fallback to team metadata
if user_api_key_dict.team_metadata:
@ -1032,20 +1042,9 @@ def get_key_model_tpm_limit(
3. Team metadata (model_tpm_limit)
4. Deployment default_api_key_tpm_limit (when model_name is provided)
"""
# 1. Check key metadata first (takes priority)
if user_api_key_dict.metadata:
result: Final = user_api_key_dict.metadata.get("model_tpm_limit")
if result:
return result
# 2. Check model_max_budget (iterate per-model like RPM does)
if user_api_key_dict.model_max_budget:
model_tpm_limit: Final[dict[str, int]] = {}
for model, budget in user_api_key_dict.model_max_budget.items():
if isinstance(budget, dict) and budget.get("tpm_limit") is not None:
model_tpm_limit[model] = budget["tpm_limit"]
if model_tpm_limit:
return model_tpm_limit
key_own_limit: Final = get_key_own_model_rate_limit(user_api_key_dict, "model_tpm_limit")
if key_own_limit is not None:
return key_own_limit
# 3. Fallback to team metadata
if user_api_key_dict.team_metadata:
@ -1967,6 +1966,11 @@ def request_dispatched_to_pass_through_endpoint(request: Request | None) -> bool
return getattr(endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, False) is True
def request_dispatched_to_provider_pass_through(request: Request) -> bool:
"""Built-in provider pass-through handlers (``/anthropic/{endpoint:path}``, ...) bind ``endpoint``."""
return "endpoint" in request.path_params
def get_model_from_request(
request_data: dict,
route: str,
@ -2040,6 +2044,12 @@ def get_model_from_request(
azure_model: Final = _router_model_from_azure_route(route, llm_router)
return model if azure_model is None else azure_model
if route.lower().startswith("/nvidia_nim/"):
nvidia_nim_model: Final = (
nvidia_nim_model_group_in_path(route, llm_router.get_model_list()) if llm_router else None
)
return model if nvidia_nim_model is None else nvidia_nim_model
return model

View file

@ -26,11 +26,13 @@ from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.caching.redis_cache import RedisCache
from litellm.constants import (
CLIENT_REQUESTED_MODEL_SCOPE_KEY,
GLOBAL_PROXY_SPEND_CACHE_KEY,
INVALID_VIRTUAL_KEY_ERROR_MARKER,
INVALID_VIRTUAL_KEY_ERROR_MESSAGE,
LITELLM_PROXY_BUDGET_NAME,
LITELLM_PROXY_MASTER_KEY_ALIAS,
MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY,
)
from litellm.integrations.otel.model.config import is_otel_v2_enabled
from litellm.integrations.otel.runtime import phase_span, seed_request_identity
@ -76,6 +78,8 @@ from litellm.proxy.auth.auth_utils import (
iter_request_fallback_targets,
normalize_request_route,
pre_db_read_auth_checks,
request_dispatched_to_pass_through_endpoint,
request_dispatched_to_provider_pass_through,
route_in_additonal_public_routes,
)
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
@ -103,6 +107,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
_safe_set_request_parsed_body,
populate_request_with_path_params,
read_raw_json_body,
rewrite_request_model,
)
from litellm.proxy.common_utils.model_listing_utils import claude_code_requested_group
from litellm.proxy.common_utils.realtime_utils import _realtime_request_body
@ -124,6 +129,7 @@ from litellm.proxy.utils import (
normalize_route_for_root_path,
)
from litellm.repositories.table_repositories import TeamMembershipRepository
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.secret_managers.main import get_secret_bool
from litellm.types.services import ServiceTypes
@ -235,11 +241,45 @@ async def _normalize_claude_model(
request.scope[_CLAUDE_MODEL_NORMALIZED] = True
if source is None:
return
request_data["model"] = source
_safe_set_request_parsed_body(request=request, parsed_body=request_data)
if request is not None:
request._json = request_data
request._body = orjson.dumps(request_data)
rewrite_request_model(request_data, request, source)
async def _resolve_router_settings_model_group_alias(
request_data: dict[str, object], # mutable-ok: the request body is rewritten in place for every downstream reader
valid_token: UserAPIKeyAuth,
request: Request | None,
route: str,
) -> None:
"""Rewrite the requested model through the key's or team's ``router_settings.model_group_alias``
before the allowlist checks, so they authorize the model group the request is routed to.
"""
from litellm.proxy.proxy_server import llm_router, prisma_client, proxy_config, proxy_logging_obj
if request is None or llm_router is None or not RouteChecks.is_llm_api_route(route=route):
return
if request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is True:
return
request.scope[MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY] = True
if request_dispatched_to_pass_through_endpoint(request) or request_dispatched_to_provider_pass_through(request):
return
requested: Final = request_data.get("model")
if not isinstance(requested, str) or await read_raw_json_body(request=request) is None:
return
settings: Final = await proxy_config.get_hierarchical_router_settings(
user_api_key_dict=valid_token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
)
if not isinstance(settings, Mapping):
return
target: Final = resolve_model_group_alias(settings.get("model_group_alias"), requested)
if target is None or target == requested:
return
verbose_proxy_logger.debug(
"router_settings.model_group_alias resolved %s -> %s before auth",
requested.replace("\r", "").replace("\n", ""),
target.replace("\r", "").replace("\n", ""),
)
request.scope.setdefault(CLIENT_REQUESTED_MODEL_SCOPE_KEY, requested)
rewrite_request_model(request_data, request, target)
def _get_model_names_for_budget_checks(
@ -269,6 +309,17 @@ class _TokenTeamModels(Protocol):
def team_models(self) -> list[str]: ...
class _RawCacheRead(Protocol):
async def async_get_cache(self, *, key: str) -> object: ...
def _raw_cache(cache: _RawCacheRead) -> _RawCacheRead:
"""View an untyped cache object's ``async_get_cache`` as returning ``object``
instead of ``Any``, so a caller can ``isinstance``-narrow it without paying
the ``reportAny`` cost of the underlying (unannotated) cache implementation."""
return cache
def _token_team_models(valid_token: _TokenTeamModels) -> list[str]:
return valid_token.team_models
@ -842,6 +893,7 @@ class _PendingAutoRegister(NamedTuple):
claim_field: str
claim_value: str
cache_key: str
jwt_issuer: str | None = None
async def _auto_register_jwt_mapping(
@ -853,6 +905,7 @@ async def _auto_register_jwt_mapping(
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging,
cache_key: str,
jwt_issuer: str | None = None,
team_id: str | None = None,
user_id: str | None = None,
org_id: str | None = None,
@ -905,6 +958,7 @@ async def _auto_register_jwt_mapping(
try:
await prisma_client.db.litellm_jwtkeymapping.create(
data={
"jwt_issuer": jwt_issuer or "",
"jwt_claim_name": virtual_key_claim_field,
"jwt_claim_value": claim_value,
"token": token_hash,
@ -939,6 +993,7 @@ async def _auto_register_jwt_mapping(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=claim_value,
prisma_client=prisma_client,
jwt_issuer=jwt_issuer,
)
if token_hash is None:
# The winner's mapping vanished between the unique-constraint
@ -983,6 +1038,43 @@ async def _auto_register_jwt_mapping(
return auto_registered_key
async def _lookup_jwt_mapping_token_hash(
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
virtual_key_claim_field: str,
claim_value: str,
normalized_issuer: str | None,
cache_key: str,
ttl: float,
) -> str | None:
issuer_scoped: Final = await get_jwt_key_mapping_object(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=claim_value,
prisma_client=prisma_client,
jwt_issuer=normalized_issuer,
)
if issuer_scoped is not None:
await user_api_key_cache.async_set_cache(key=cache_key, value=issuer_scoped, ttl=ttl)
return issuer_scoped
if normalized_issuer is None:
return None
# Another issuer may have already resolved (and cached) this same
# global mapping -- check its cache entry before re-querying the DB.
global_cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, claim_value)
cached_global: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=global_cache_key)
if isinstance(cached_global, str) and cached_global != "__NO_MAPPING__":
return cached_global
global_row: Final = await get_jwt_key_mapping_object(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=claim_value,
prisma_client=prisma_client,
jwt_issuer=None,
)
if global_row is not None:
await user_api_key_cache.async_set_cache(key=global_cache_key, value=global_row, ttl=ttl)
return global_row
async def _resolve_jwt_to_virtual_key(
jwt_claims: dict,
jwt_handler: JWTHandler,
@ -1041,7 +1133,7 @@ async def _resolve_jwt_to_virtual_key(
)
return None
cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value))
cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value), normalized_issuer)
raw_cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key)
sentinel_written_by_this_policy: Final = behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER
cached_mapping: Final = (
@ -1081,6 +1173,7 @@ async def _resolve_jwt_to_virtual_key(
claim_field=virtual_key_claim_field,
claim_value=str(claim_value),
cache_key=cache_key,
jwt_issuer=normalized_issuer,
)
return None
elif cached_mapping is not None:
@ -1094,21 +1187,30 @@ async def _resolve_jwt_to_virtual_key(
)
# Resolve the mapping from DB, or treat prisma_client=None as a definitive
# miss (no DB → no mapping can exist → apply no-match policy below).
token_hash: str | None = None
if prisma_client is not None:
token_hash = await get_jwt_key_mapping_object(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=str(claim_value),
# miss (no DB → no mapping can exist → apply no-match policy below). An
# issuer-scoped row wins; falling back to the global (no-issuer) row keeps
# mappings created before issuer scoping existed working for every issuer.
# Each tier is cached under ITS OWN key (the global tier under the
# issuer-less cache key, not under `cache_key`/this issuer's key) so that
# updating or deleting either row invalidates exactly the cache entries it
# can affect. Caching a global-row hit under the requesting issuer's key
# would leave every OTHER issuer that had fallen back to that same global
# mapping serving its stale token until TTL after the row changes.
token_hash: Final = (
await _lookup_jwt_mapping_token_hash(
prisma_client=prisma_client,
)
if token_hash is not None:
await user_api_key_cache.async_set_cache(
key=cache_key,
value=token_hash,
user_api_key_cache=user_api_key_cache,
virtual_key_claim_field=virtual_key_claim_field,
claim_value=str(claim_value),
normalized_issuer=normalized_issuer,
cache_key=cache_key,
ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl,
)
if prisma_client is not None
else None
)
if token_hash is not None:
return IdentityStore.key_from_principal(
await IdentityStore(
prisma_client,
@ -1149,6 +1251,7 @@ async def _resolve_jwt_to_virtual_key(
claim_field=virtual_key_claim_field,
claim_value=str(claim_value),
cache_key=cache_key,
jwt_issuer=normalized_issuer,
)
# FALLBACK_TEAM_MAPPING (default): cache the miss and return None so the
@ -1641,6 +1744,7 @@ async def _user_api_key_auth_builder(
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
cache_key=pending_auto_register.cache_key,
jwt_issuer=pending_auto_register.jwt_issuer,
team_id=team_id,
user_id=user_id,
org_id=org_id,
@ -2926,6 +3030,7 @@ async def _authorize_authenticated_request(
## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request)
await _normalize_claude_model(request_data, user_api_key_auth_obj, request, route)
await _resolve_router_settings_model_group_alias(request_data, user_api_key_auth_obj, request, route)
# Single authorization point. Builder paths MUST NOT call common_checks.
# Route through the same exception handler the builder uses so
@ -3312,6 +3417,7 @@ async def _enforce_key_and_fallback_model_access(
Not included in common_checks common_checks enforces team/user/project model access only.
"""
await _normalize_claude_model(request_data, valid_token, request, route)
await _resolve_router_settings_model_group_alias(request_data, valid_token, request, route)
config: Final = valid_token.config
if config != {}:

View file

@ -490,7 +490,7 @@ lite codex exec "summarize the repo"
Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process.
The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, so the proxy lists every other group to Claude Code as `claude-router-<UTF-8 hex of the group name>` and marks a group whose input window reaches 1M with `[1m]`, and a request on such an id is served by the group. Older Claude Code versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway.
The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, so the proxy lists every other group to Claude Code as `claude-router-<UTF-8 hex of the group name>` and marks a group whose input window reaches 1M with `[1m]`, and a request on such an id is served by the group. Older Claude Code versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway. Codex gets the same list as a catalog file, `$CODEX_HOME/litellm-models.json` (default `~/.codex/`), passed as `-c model_catalog_json=<path>` so `/model` lists exactly the proxy's chat models; a proxy model the installed Codex already knows (`gpt-5.5`, say) keeps that Codex's own entry, reasoning levels and prompt included, and only its place in the picker comes from the proxy, while a model Codex does not know gets the plain entry Codex uses for an unknown `-m` slug. Before launching, `lite codex` asks the installed Codex for its own list and then has it read the written file back (both through `codex debug models`), and when the fetch, either of those or the write fails (Codex releases older than 0.130 have no such command) it says so on stderr and launches with Codex's built-in catalog, leaving a rejected file in place.
pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/<id>`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. Your own flags come after the injected pin, so `lite pi --model litellm/<other-id>` wins, and inside the TUI the `/model` picker lists every synced litellm model.

View file

@ -4,15 +4,16 @@ import re
import shutil
import subprocess
import sys
import tempfile
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final, TypeAlias
from typing import Final, Literal, TypeAlias
import click
import requests
from pydantic import BaseModel, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from .auth import CliContextObj, context_secret_vault, get_stored_api_key, login
from .claude_settings import ClaudeSettingsError, install_statusline_script
@ -65,6 +66,10 @@ _INSTALL_DOCS: Final[dict[str, str]] = {
_HIDDEN_AGENTS: Final = frozenset({"pi"})
CODEX_PROXY_PROVIDER: Final = "litellm"
CODEX_HOME_ENV: Final = "CODEX_HOME"
CODEX_MODEL_CATALOG_FILENAME: Final = "litellm-models.json"
_CODEX_BASE_INSTRUCTIONS_PATH: Final = Path(__file__).with_name("codex_base_instructions.md")
_CODEX_PREFLIGHT_TIMEOUT_SECONDS: Final = 10.0
class AgentRunError(Exception):
@ -252,7 +257,7 @@ def agent_launch_args(command: str, base_url: str) -> list[str]:
class ListedModel(BaseModel):
"""The fields of a /v1/models entry that an OpenCode model entry is built from."""
"""The fields of a /v1/models entry that an OpenCode or Codex model entry is built from."""
id: str
mode: str | None = None
@ -265,7 +270,7 @@ class _ModelListing(BaseModel):
_MODEL_LISTING: Final = TypeAdapter(_ModelListing)
_OPENCODE_CHAT_MODES: Final[frozenset[str]] = frozenset({"chat", "responses"})
_CHAT_MODES: Final[frozenset[str]] = frozenset({"chat", "responses"})
_NO_EXTRA_ENV: Final[Mapping[str, str]] = MappingProxyType({})
@ -274,6 +279,40 @@ class ModelSyncSkipped:
reason: str
@dataclass(frozen=True, slots=True)
class ModelSyncArgs:
"""CLI args, placed before the user's own, that hand an agent the synced model list."""
args: tuple[str, ...]
ModelSyncResult: TypeAlias = Mapping[str, str] | ModelSyncArgs | ModelSyncSkipped
def _chat_models(models: Sequence[ListedModel]) -> tuple[ListedModel, ...]:
return tuple(m for m in models if m.mode is None or m.mode in _CHAT_MODES)
def _fetch_model_listing(
base_url: str,
api_key: str,
*,
get: Callable[..., requests.Response],
) -> tuple[ListedModel, ...] | ModelSyncSkipped:
url: Final = base_url.rstrip("/") + "/v1/models"
try:
resp: Final = get(url, headers=MappingProxyType({"Authorization": f"Bearer {api_key}"}), timeout=10)
except requests.RequestException as e:
return ModelSyncSkipped(f"could not reach {url}: {e}")
if resp.status_code != 200:
return ModelSyncSkipped(f"{url} returned HTTP {resp.status_code}")
try:
listing: Final = _MODEL_LISTING.validate_json(resp.content)
except ValidationError:
return ModelSyncSkipped(f"{url} returned an unexpected body")
return listing.data
class _OpenCodeLimit(BaseModel):
context: int
output: int
@ -317,7 +356,7 @@ def opencode_provider_config(base_url: str, models: Sequence[ListedModel]) -> st
it never lands in the config text. OpenCode merges this inline config over
the user's own files, leaving unrelated keys and providers untouched.
"""
chat_models: Final = tuple(m for m in models if m.mode is None or m.mode in _OPENCODE_CHAT_MODES)
chat_models: Final = _chat_models(models)
provider: Final = _OpenCodeProvider(
npm=OPENCODE_PROVIDER_NPM,
name=OPENCODE_PROVIDER_NAME,
@ -347,40 +386,269 @@ def opencode_model_sync_env(
"""
if OPENCODE_CONFIG_CONTENT_ENV in base_env:
return ModelSyncSkipped(f"{OPENCODE_CONFIG_CONTENT_ENV} is already set")
url: Final = base_url.rstrip("/") + "/v1/models"
listing: Final = _fetch_model_listing(base_url, api_key, get=get)
if isinstance(listing, ModelSyncSkipped):
return listing
return MappingProxyType({OPENCODE_CONFIG_CONTENT_ENV: opencode_provider_config(base_url, listing)})
class _CodexTruncationPolicy(BaseModel):
mode: Literal["bytes"] = "bytes"
limit: int = 10_000
class _CodexModel(BaseModel):
"""One `ModelInfo` entry of a Codex model catalog for a model the installed Codex does not know.
Every field that some Codex release since `model_catalog_json` appeared
(0.105.0) deserializes without a default is spelled out here, so one catalog
parses on all of them; the values match the fallback metadata Codex uses for
a model slug it does not know, so picking such a proxy model behaves the
same as `codex -m` did.
"""
slug: str
display_name: str
description: None = None
supported_reasoning_levels: tuple[()] = ()
shell_type: Literal["unified_exec"] = "unified_exec"
visibility: Literal["list"] = "list"
supported_in_api: Literal[True] = True
priority: int
availability_nux: None = None
upgrade: None = None
support_verbosity: Literal[False] = False
supports_reasoning_summaries: Literal[False] = False
supports_parallel_tool_calls: Literal[False] = False
default_verbosity: None = None
apply_patch_tool_type: None = None
truncation_policy: _CodexTruncationPolicy = _CodexTruncationPolicy()
experimental_supported_tools: tuple[()] = ()
context_window: int | None
base_instructions: str
class _StockCodexUpgrade(BaseModel):
model_config = ConfigDict(extra="allow")
model: str
class _StockCodexModel(BaseModel):
"""One `ModelInfo` entry as the installed Codex prints it from `codex debug models`.
Only the fields the sync rewrites are named; everything else that release
knows about the model (its reasoning levels, prompt, tool support) rides
along untouched, whatever the release's schema.
"""
model_config = ConfigDict(extra="allow")
slug: str
priority: int
visibility: str
supported_in_api: bool = True
upgrade: _StockCodexUpgrade | None = None
class _StockCodexCatalog(BaseModel):
models: tuple[_StockCodexModel, ...]
class _CodexCatalog(BaseModel):
models: tuple[_CodexModel | _StockCodexModel, ...]
def _codex_catalog_entry(
priority: int,
listed: ListedModel,
stock: _StockCodexModel | None,
served: frozenset[str],
instructions: str,
) -> _CodexModel | _StockCodexModel:
if stock is None:
return _CodexModel(
slug=listed.id,
display_name=listed.id,
priority=priority,
context_window=listed.max_input_tokens,
base_instructions=instructions,
)
upgrade: Final = stock.upgrade if stock.upgrade is not None and stock.upgrade.model in served else None
return stock.model_copy(
update={"priority": priority, "visibility": "list", "supported_in_api": True, "upgrade": upgrade}
)
def codex_model_catalog(
models: Sequence[ListedModel], stock: Sequence[_StockCodexModel], instructions: str
) -> str | None:
"""The `model_catalog_json` body listing the proxy's chat models, or None if there are none.
Codex refuses an empty catalog, hence None instead of `{"models": []}`.
Passing a catalog replaces Codex's built-in one, so a proxy model the
installed Codex knows keeps that Codex's own entry and the proxy only
decides its place in the picker: the listing orders it, lists it even when
Codex hides it or keeps it off the API, and keeps Codex's upgrade nudge only
when the model it points at is served too. A model Codex does not know gets the fallback
entry, with the same base instructions Codex itself uses so the agent never
runs without a system prompt.
"""
chat_models: Final = _chat_models(models)
if not chat_models:
return None
served: Final = frozenset(m.id for m in chat_models)
known: Final = MappingProxyType({m.slug: m for m in stock})
catalog: Final = _CodexCatalog(
models=tuple(
_codex_catalog_entry(index, m, known.get(m.id), served, instructions) for index, m in enumerate(chat_models)
)
)
return catalog.model_dump_json()
def codex_model_catalog_path(env: Mapping[str, str], *, home: Callable[[], Path] = Path.home) -> Path:
override: Final = env.get(CODEX_HOME_ENV)
root: Final = Path(override) if override else home() / ".codex"
return root / CODEX_MODEL_CATALOG_FILENAME
def _replace_file(path: Path, text: str) -> None:
with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as tmp:
_ = tmp.write(text)
try:
resp: Final = get(url, headers=MappingProxyType({"Authorization": f"Bearer {api_key}"}), timeout=10)
except requests.RequestException as e:
return ModelSyncSkipped(f"could not reach {url}: {e}")
if resp.status_code != 200:
return ModelSyncSkipped(f"{url} returned HTTP {resp.status_code}")
os.replace(tmp.name, path)
except OSError:
Path(tmp.name).unlink(missing_ok=True)
raise
def _codex_debug_models(
binary: str,
args: Sequence[str],
env: Mapping[str, str],
*,
run: Callable[..., subprocess.CompletedProcess[str]],
) -> str | ModelSyncSkipped:
"""What `codex debug models` prints with `args` in front, or why the installed Codex could not run it.
The command prints the catalog Codex would launch with, without touching
the network, so it lists the installed Codex's own models and parses a
catalog override the way a launch does. Releases before 0.130.0 have no
such command and are reported the same way. A batch shim goes through
cmd.exe exactly as the launch will.
"""
name: Final = os.path.basename(binary)
command: Final = _windows_command(binary, (binary, *args, "debug", "models"))
try:
listing: Final = _MODEL_LISTING.validate_json(resp.content)
except ValidationError:
return ModelSyncSkipped(f"{url} returned an unexpected body")
return MappingProxyType({OPENCODE_CONFIG_CONTENT_ENV: opencode_provider_config(base_url, listing.data)})
completed: Final = run(
command,
env=dict(env),
stdin=subprocess.DEVNULL,
capture_output=True,
encoding="utf-8",
timeout=_CODEX_PREFLIGHT_TIMEOUT_SECONDS,
)
except (OSError, subprocess.TimeoutExpired) as e:
return ModelSyncSkipped(f"`{name} debug models` failed: {e}")
if completed.returncode == 0:
return completed.stdout
lines: Final = completed.stderr.strip().splitlines()
detail: Final = lines[0] if lines else "no output"
return ModelSyncSkipped(f"`{name} debug models` exited {completed.returncode}: {detail}")
def _stock_codex_models(
binary: str, env: Mapping[str, str], *, run: Callable[..., subprocess.CompletedProcess[str]]
) -> tuple[_StockCodexModel, ...] | ModelSyncSkipped:
printed: Final = _codex_debug_models(binary, (), env, run=run)
if isinstance(printed, ModelSyncSkipped):
return printed
try:
return _StockCodexCatalog.model_validate_json(printed).models
except ValidationError as e:
name: Final = os.path.basename(binary)
return ModelSyncSkipped(f"`{name} debug models` printed no model catalog: {e.errors()[0]['msg']}")
def codex_model_sync_args(
base_env: Mapping[str, str],
base_url: str,
api_key: str,
*,
binary: str = "codex",
get: Callable[..., requests.Response] = requests.get,
run: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
home: Callable[[], Path] = Path.home,
instructions_path: Path = _CODEX_BASE_INSTRUCTIONS_PATH,
) -> ModelSyncArgs | ModelSyncSkipped:
"""`-c model_catalog_json=...` pointing Codex at the proxy's model list, or why it was skipped.
Codex has no env or inline equivalent of OPENCODE_CONFIG_CONTENT: the catalog
must be a file, so it is written under $CODEX_HOME (default ~/.codex) and
atomically replaced on every launch. The Codex at `binary` first lists its
own models, so the ones the proxy serves keep that Codex's entries, and then
reads the file back once before it is handed over. The key never lands in
the file. A failed fetch, read, listing, write or read-back is reported
rather than raised: Codex still launches with its built-in catalog and takes
a proxy model by name via -m, and a rejected file stays on disk to be looked
at.
"""
listing: Final = _fetch_model_listing(base_url, api_key, get=get)
if isinstance(listing, ModelSyncSkipped):
return listing
try:
instructions: Final = instructions_path.read_text(encoding="utf-8")
except OSError as e:
return ModelSyncSkipped(f"could not read {instructions_path}: {e}")
path: Final = codex_model_catalog_path(base_env, home=home)
try:
path.parent.mkdir(parents=True, exist_ok=True)
except OSError as e:
return ModelSyncSkipped(f"could not write {path}: {e}")
stock: Final = _stock_codex_models(binary, base_env, run=run)
if isinstance(stock, ModelSyncSkipped):
return stock
catalog: Final = codex_model_catalog(listing, stock, instructions)
if catalog is None:
return ModelSyncSkipped(f"{base_url.rstrip('/')}/v1/models lists no chat models")
try:
_replace_file(path, catalog)
except OSError as e:
return ModelSyncSkipped(f"could not write {path}: {e}")
override: Final = f"model_catalog_json={json.dumps(str(path))}"
read_back: Final = _codex_debug_models(binary, ("-c", override), base_env, run=run)
if isinstance(read_back, ModelSyncSkipped):
return read_back
return ModelSyncArgs(("-c", override))
def agent_model_sync_env(
command: str,
binary: str,
base_env: Mapping[str, str],
base_url: str,
api_key: str,
skip_verify: bool,
*,
get: Callable[..., requests.Response] = requests.get,
) -> Mapping[str, str] | ModelSyncSkipped:
"""Extra env an agent needs to see the proxy's model list.
run: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
) -> ModelSyncResult:
"""Extra env or args an agent needs to see the proxy's model list.
Only OpenCode needs one: Claude Code discovers models through
CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY and Codex takes the model by name.
skip_verify means the caller wants no pre-launch proxy call at all, so the
listing is skipped too rather than hanging on an offline proxy.
binary is the resolved path the launch will run (`codex.cmd` on a Windows
npm install). OpenCode takes the list as env, Codex as a `-c` override that
binary has read back first; Claude Code discovers models itself through
CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY. skip_verify means the caller
wants no pre-launch proxy call at all, so the listing is skipped too rather
than hanging on an offline proxy.
"""
if os.path.basename(command) != "opencode":
agent: Final = os.path.splitext(os.path.basename(binary))[0]
if agent not in ("opencode", "codex"):
return _NO_EXTRA_ENV
if skip_verify:
return ModelSyncSkipped(f"{_SKIP_VERIFY_FLAG} was passed")
if agent == "codex":
return codex_model_sync_args(base_env, base_url, api_key, binary=binary, get=get, run=run)
return opencode_model_sync_env(base_env, base_url, api_key, get=get)
@ -508,9 +776,7 @@ def run_agent(
base_env: Mapping[str, str] | None = None,
which: Callable[[str], str | None] = shutil.which,
verify: Callable[[str, str], None] = verify_proxy_key,
sync_models: Callable[[str, Mapping[str, str], str, str, bool], Mapping[str, str] | ModelSyncSkipped] = (
agent_model_sync_env
),
sync_models: Callable[[str, Mapping[str, str], str, str, bool], ModelSyncResult] = agent_model_sync_env,
warn: Callable[[str], None] = _warn,
launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _hand_off,
reattach_terminal: Callable[[], None] | None = None,
@ -537,7 +803,7 @@ def run_agent(
verify(base_url, api_key)
env_before_sync: Final = base_env if base_env is not None else os.environ
synced: Final = sync_models(command[0], env_before_sync, base_url, api_key, skip_verify)
synced: Final = sync_models(binary, env_before_sync, base_url, api_key, skip_verify)
if isinstance(synced, ModelSyncSkipped):
warn(f"litellm: not syncing {display_name} models from the proxy: {synced.reason}")
@ -547,10 +813,11 @@ def run_agent(
env: Final = MappingProxyType(
{
**build_agent_env(env_before_sync, base_url, api_key, profiles),
**(_NO_EXTRA_ENV if isinstance(synced, ModelSyncSkipped) else synced),
**(synced if isinstance(synced, Mapping) else _NO_EXTRA_ENV),
}
)
extra_args: Final = (*agent_launch_args(command[0], base_url), *prepared_args)
synced_args: Final = synced.args if isinstance(synced, ModelSyncArgs) else ()
extra_args: Final = (*agent_launch_args(command[0], base_url), *synced_args, *prepared_args)
if reattach_terminal is not None:
reattach_terminal()
launcher(binary, [command[0], *extra_args, *command[1:]], env)

View file

@ -0,0 +1,275 @@
You are a coding agent running in the Codex CLI, a terminal-based coding assistant. Codex CLI is an open source project led by OpenAI. You are expected to be precise, safe, and helpful.
Your capabilities:
- Receive user prompts and other context provided by the harness, such as files in the workspace.
- Communicate with the user by streaming thinking & responses, and by making & updating plans.
- Emit function calls to run terminal commands and apply patches. Depending on how this specific run is configured, you can request that these function calls be escalated to the user for approval before running. More on this in the "Sandbox and approvals" section.
Within this context, Codex refers to the open-source agentic coding interface (not the old Codex language model built by OpenAI).
# How you work
## Personality
Your default personality and tone is concise, direct, and friendly. You communicate efficiently, always keeping the user clearly informed about ongoing actions without unnecessary detail. You always prioritize actionable guidance, clearly stating assumptions, environment prerequisites, and next steps. Unless explicitly asked, you avoid excessively verbose explanations about your work.
# AGENTS.md spec
- Repos often contain AGENTS.md files. These files can appear anywhere within the repository.
- These files are a way for humans to give you (the agent) instructions or tips for working within the container.
- Some examples might be: coding conventions, info about how code is organized, or instructions for how to run or test code.
- Instructions in AGENTS.md files:
- The scope of an AGENTS.md file is the entire directory tree rooted at the folder that contains it.
- For every file you touch in the final patch, you must obey instructions in any AGENTS.md file whose scope includes that file.
- Instructions about code style, structure, naming, etc. apply only to code within the AGENTS.md file's scope, unless the file states otherwise.
- More-deeply-nested AGENTS.md files take precedence in the case of conflicting instructions.
- Direct system/developer/user instructions (as part of a prompt) take precedence over AGENTS.md instructions.
- The contents of the AGENTS.md file at the root of the repo and any directories from the CWD up to the root are included with the developer message and don't need to be re-read. When working in a subdirectory of CWD, or a directory outside the CWD, check for any AGENTS.md files that may be applicable.
## Responsiveness
### Preamble messages
Before making tool calls, send a brief preamble to the user explaining what youre about to do. When sending preamble messages, follow these principles and examples:
- **Logically group related actions**: if youre about to run several related commands, describe them together in one preamble rather than sending a separate note for each.
- **Keep it concise**: be no more than 1-2 sentences, focused on immediate, tangible next steps. (812 words for quick updates).
- **Build on prior context**: if this is not your first tool call, use the preamble message to connect the dots with whats been done so far and create a sense of momentum and clarity for the user to understand your next actions.
- **Keep your tone light, friendly and curious**: add small touches of personality in preambles feel collaborative and engaging.
- **Exception**: Avoid adding a preamble for every trivial read (e.g., `cat` a single file) unless its part of a larger grouped action.
**Examples:**
- “Ive explored the repo; now checking the API route definitions.”
- “Next, Ill patch the config and update the related tests.”
- “Im about to scaffold the CLI commands and helper functions.”
- “Ok cool, so Ive wrapped my head around the repo. Now digging into the API routes.”
- “Configs looking tidy. Next up is patching helpers to keep things in sync.”
- “Finished poking at the DB gateway. I will now chase down error handling.”
- “Alright, build pipeline order is interesting. Checking how it reports failures.”
- “Spotted a clever caching util; now hunting where it gets used.”
## Planning
You have access to an `update_plan` tool which tracks steps and progress and renders them to the user. Using the tool helps demonstrate that you've understood the task and convey how you're approaching it. Plans can help to make complex, ambiguous, or multi-phase work clearer and more collaborative for the user. A good plan should break the task into meaningful, logically ordered steps that are easy to verify as you go.
Note that plans are not for padding out simple work with filler steps or stating the obvious. The content of your plan should not involve doing anything that you aren't capable of doing (i.e. don't try to test things that you can't test). Do not use plans for simple or single-step queries that you can just do or answer immediately.
Do not repeat the full contents of the plan after an `update_plan` call — the harness already displays it. Instead, summarize the change made and highlight any important context or next step.
Before running a command, consider whether or not you have completed the previous step, and make sure to mark it as completed before moving on to the next step. It may be the case that you complete all steps in your plan after a single pass of implementation. If this is the case, you can simply mark all the planned steps as completed. Sometimes, you may need to change plans in the middle of a task: call `update_plan` with the updated plan and make sure to provide an `explanation` of the rationale when doing so.
Use a plan when:
- The task is non-trivial and will require multiple actions over a long time horizon.
- There are logical phases or dependencies where sequencing matters.
- The work has ambiguity that benefits from outlining high-level goals.
- You want intermediate checkpoints for feedback and validation.
- When the user asked you to do more than one thing in a single prompt
- The user has asked you to use the plan tool (aka "TODOs")
- You generate additional steps while working, and plan to do them before yielding to the user
### Examples
**High-quality plans**
Example 1:
1. Add CLI entry with file args
2. Parse Markdown via CommonMark library
3. Apply semantic HTML template
4. Handle code blocks, images, links
5. Add error handling for invalid files
Example 2:
1. Define CSS variables for colors
2. Add toggle with localStorage state
3. Refactor components to use variables
4. Verify all views for readability
5. Add smooth theme-change transition
Example 3:
1. Set up Node.js + WebSocket server
2. Add join/leave broadcast events
3. Implement messaging with timestamps
4. Add usernames + mention highlighting
5. Persist messages in lightweight DB
6. Add typing indicators + unread count
**Low-quality plans**
Example 1:
1. Create CLI tool
2. Add Markdown parser
3. Convert to HTML
Example 2:
1. Add dark mode toggle
2. Save preference
3. Make styles look good
Example 3:
1. Create single-file HTML game
2. Run quick sanity check
3. Summarize usage instructions
If you need to write a plan, only write high quality plans, not low quality ones.
## Task execution
You are a coding agent. Please keep going until the query is completely resolved, before ending your turn and yielding back to the user. Only terminate your turn when you are sure that the problem is solved. Autonomously resolve the query to the best of your ability, using the tools available to you, before coming back to the user. Do NOT guess or make up an answer.
You MUST adhere to the following criteria when solving queries:
- Working on the repo(s) in the current environment is allowed, even if they are proprietary.
- Analyzing code for vulnerabilities is allowed.
- Showing user code and tool call details is allowed.
- Use the `apply_patch` tool to edit files (NEVER try `applypatch` or `apply-patch`, only `apply_patch`): {"command":["apply_patch","*** Begin Patch\\n*** Update File: path/to/file.py\\n@@ def example():\\n- pass\\n+ return 123\\n*** End Patch"]}
If completing the user's task requires writing or modifying files, your code and final answer should follow these coding guidelines, though user instructions (i.e. AGENTS.md) may override these guidelines:
- Fix the problem at the root cause rather than applying surface-level patches, when possible.
- Avoid unneeded complexity in your solution.
- Do not attempt to fix unrelated bugs or broken tests. It is not your responsibility to fix them. (You may mention them to the user in your final message though.)
- Update documentation as necessary.
- Keep changes consistent with the style of the existing codebase. Changes should be minimal and focused on the task.
- Use `git log` and `git blame` to search the history of the codebase if additional context is required.
- NEVER add copyright or license headers unless specifically requested.
- Do not waste tokens by re-reading files after calling `apply_patch` on them. The tool call will fail if it didn't work. The same goes for making folders, deleting folders, etc.
- Do not `git commit` your changes or create new git branches unless explicitly requested.
- Do not add inline comments within code unless explicitly requested.
- Do not use one-letter variable names unless explicitly requested.
- NEVER output inline citations like "【F:README.md†L5-L14】" in your outputs. The CLI is not able to render these so they will just be broken in the UI. Instead, if you output valid filepaths, users will be able to click on them to open the files in their editor.
## Validating your work
If the codebase has tests or the ability to build or run, consider using them to verify that your work is complete.
When testing, your philosophy should be to start as specific as possible to the code you changed so that you can catch issues efficiently, then make your way to broader tests as you build confidence. If there's no test for the code you changed, and if the adjacent patterns in the codebases show that there's a logical place for you to add a test, you may do so. However, do not add tests to codebases with no tests.
Similarly, once you're confident in correctness, you can suggest or use formatting commands to ensure that your code is well formatted. If there are issues you can iterate up to 3 times to get formatting right, but if you still can't manage it's better to save the user time and present them a correct solution where you call out the formatting in your final message. If the codebase does not have a formatter configured, do not add one.
For all of testing, running, building, and formatting, do not attempt to fix unrelated bugs. It is not your responsibility to fix them. (You may mention them to the user in your final message though.)
Be mindful of whether to run validation commands proactively. In the absence of behavioral guidance:
- When running in the non-interactive approval mode **never**, proactively run tests, lint and do whatever you need to ensure you've completed the task.
- When working in interactive approval modes like **untrusted**, or **on-request**, hold off on running tests or lint commands until the user is ready for you to finalize your output, because these commands take time to run and slow down iteration. Instead suggest what you want to do next, and let the user confirm first.
- When working on test-related tasks, such as adding tests, fixing tests, or reproducing a bug to verify behavior, you may proactively run tests regardless of approval mode. Use your judgement to decide whether this is a test-related task.
## Ambition vs. precision
For tasks that have no prior context (i.e. the user is starting something brand new), you should feel free to be ambitious and demonstrate creativity with your implementation.
If you're operating in an existing codebase, you should make sure you do exactly what the user asks with surgical precision. Treat the surrounding codebase with respect, and don't overstep (i.e. changing filenames or variables unnecessarily). You should balance being sufficiently ambitious and proactive when completing tasks of this nature.
You should use judicious initiative to decide on the right level of detail and complexity to deliver based on the user's needs. This means showing good judgment that you're capable of doing the right extras without gold-plating. This might be demonstrated by high-value, creative touches when scope of the task is vague; while being surgical and targeted when scope is tightly specified.
## Sharing progress updates
For especially longer tasks that you work on (i.e. requiring many tool calls, or a plan with multiple steps), you should provide progress updates back to the user at reasonable intervals. These updates should be structured as a concise sentence or two (no more than 8-10 words long) recapping progress so far in plain language: this update demonstrates your understanding of what needs to be done, progress so far (i.e. files explores, subtasks complete), and where you're going next.
Before doing large chunks of work that may incur latency as experienced by the user (i.e. writing a new file), you should send a concise message to the user with an update indicating what you're about to do to ensure they know what you're spending time on. Don't start editing or writing large files before informing the user what you are doing and why.
The messages you send before tool calls should describe what is immediately about to be done next in very concise language. If there was previous work done, this preamble message should also include a note about the work done so far to bring the user along.
## Presenting your work and final message
Your final message should read naturally, like an update from a concise teammate. For casual conversation, brainstorming tasks, or quick questions from the user, respond in a friendly, conversational tone. You should ask questions, suggest ideas, and adapt to the users style. If you've finished a large amount of work, when describing what you've done to the user, you should follow the final answer formatting guidelines to communicate substantive changes. You don't need to add structured formatting for one-word answers, greetings, or purely conversational exchanges.
You can skip heavy formatting for single, simple actions or confirmations. In these cases, respond in plain sentences with any relevant next step or quick option. Reserve multi-section structured responses for results that need grouping or explanation.
The user is working on the same computer as you, and has access to your work. As such there's no need to show the full contents of large files you have already written unless the user explicitly asks for them. Similarly, if you've created or modified files using `apply_patch`, there's no need to tell users to "save the file" or "copy the code into a file"—just reference the file path.
If there's something that you think you could help with as a logical next step, concisely ask the user if they want you to do so. Good examples of this are running tests, committing changes, or building out the next logical component. If theres something that you couldn't do (even with approval) but that the user might want to do (such as verifying changes by running the app), include those instructions succinctly.
Brevity is very important as a default. You should be very concise (i.e. no more than 10 lines), but can relax this requirement for tasks where additional detail and comprehensiveness is important for the user's understanding.
### Final answer structure and style guidelines
You are producing plain text that will later be styled by the CLI. Follow these rules exactly. Formatting should make results easy to scan, but not feel mechanical. Use judgment to decide how much structure adds value.
**Section Headers**
- Use only when they improve clarity — they are not mandatory for every answer.
- Choose descriptive names that fit the content
- Keep headers short (13 words) and in `**Title Case**`. Always start headers with `**` and end with `**`
- Leave no blank line before the first bullet under a header.
- Section headers should only be used where they genuinely improve scanability; avoid fragmenting the answer.
**Bullets**
- Use `-` followed by a space for every bullet.
- Merge related points when possible; avoid a bullet for every trivial detail.
- Keep bullets to one line unless breaking for clarity is unavoidable.
- Group into short lists (46 bullets) ordered by importance.
- Use consistent keyword phrasing and formatting across sections.
**Monospace**
- Wrap all commands, file paths, env vars, and code identifiers in backticks (`` `...` ``).
- Apply to inline examples and to bullet keywords if the keyword itself is a literal file/command.
- Never mix monospace and bold markers; choose one based on whether its a keyword (`**`) or inline code/path (`` ` ``).
**File References**
When referencing files in your response, make sure to include the relevant start line and always follow the below rules:
* Use inline code to make file paths clickable.
* Each reference should have a stand alone path. Even if it's the same file.
* Accepted: absolute, workspacerelative, a/ or b/ diff prefixes, or bare filename/suffix.
* Line/column (1based, optional): :line[:column] or #Lline[Ccolumn] (column defaults to 1).
* Do not use URIs like file://, vscode://, or https://.
* Do not provide range of lines
* Examples: src/app.ts, src/app.ts:42, b/server/index.js#L10, C:\repo\project\main.rs:12:5
**Structure**
- Place related bullets together; dont mix unrelated concepts in the same section.
- Order sections from general → specific → supporting info.
- For subsections (e.g., “Binaries” under “Rust Workspace”), introduce with a bolded keyword bullet, then list items under it.
- Match structure to complexity:
- Multi-part or detailed results → use clear headers and grouped bullets.
- Simple results → minimal headers, possibly just a short list or paragraph.
**Tone**
- Keep the voice collaborative and natural, like a coding partner handing off work.
- Be concise and factual — no filler or conversational commentary and avoid unnecessary repetition
- Use present tense and active voice (e.g., “Runs tests” not “This will run tests”).
- Keep descriptions self-contained; dont refer to “above” or “below”.
- Use parallel structure in lists for consistency.
**Dont**
- Dont use literal words “bold” or “monospace” in the content.
- Dont nest bullets or create deep hierarchies.
- Dont output ANSI escape codes directly — the CLI renderer applies them.
- Dont cram unrelated keywords into a single bullet; split for clarity.
- Dont let keyword lists run long — wrap or reformat for scanability.
Generally, ensure your final answers adapt their shape and depth to the request. For example, answers to code explanations should have a precise, structured explanation with code references that answer the question directly. For tasks with a simple implementation, lead with the outcome and supplement only with whats needed for clarity. Larger changes can be presented as a logical walkthrough of your approach, grouping related steps, explaining rationale where it adds value, and highlighting next actions to accelerate the user. Your answers should provide the right level of detail while being easily scannable.
For casual greetings, acknowledgements, or other one-off conversational messages that are not delivering substantive information or structured results, respond naturally without section headers or bullet formatting.
# Tool Guidelines
## Shell commands
When using the shell, you must adhere to the following guidelines:
- When searching for text or files, prefer using `rg` or `rg --files` respectively because `rg` is much faster than alternatives like `grep`. (If the `rg` command is not found, then use alternatives.)
- Do not use python scripts to attempt to output larger chunks of a file.
## `update_plan`
A tool named `update_plan` is available to you. You can use it to keep an uptodate, stepbystep plan for the task.
To create a new plan, call `update_plan` with a short list of 1sentence steps (no more than 5-7 words each) with a `status` for each step (`pending`, `in_progress`, or `completed`).
When steps have been completed, use `update_plan` to mark each finished step as `completed` and the next step you are working on as `in_progress`. There should always be exactly one `in_progress` step until everything is done. You can mark multiple items as complete in a single `update_plan` call.
If all steps are complete, ensure you call `update_plan` to mark all steps as `completed`.

View file

@ -56,6 +56,7 @@ from litellm.proxy.common_utils.callback_utils import (
get_logging_caching_headers,
get_remaining_tokens_and_requests_from_request_data,
)
from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model
from litellm.proxy.common_utils.openai_error_payload import (
attribute_of,
error_status_code,
@ -622,9 +623,9 @@ async def _resolve_per_request_model_group_alias(
holds the global config map and is shared across requests, so a per-request
map has to be applied here instead of being forwarded to the Router.
Model access was authorized against the requested group, so the target is
authorized in its own right before the rewrite; a key that may not call the
target gets the usual 403 rather than being quietly served it.
Auth already rewrote the body through this map for LLM API routes, so this is
a fallback for callers that skipped it; the target is authorized in its own
right before the rewrite, so a key that may not call it gets the usual 403.
Returns the target model group, or None when no alias applies.
"""
@ -1451,10 +1452,13 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool:
_CLIENT_DISCONNECT_DETAIL: Final = "Client disconnected the request"
def _log_llm_api_exception(e: Exception) -> None:
def _log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None:
if getattr(e, "status_code", None) == 499 and getattr(e, "detail", None) == _CLIENT_DISCONNECT_DETAIL:
verbose_proxy_logger.info(
"litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, upstream LLM request cancelled"
"litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, "
"upstream LLM request cancelled - litellm_call_id=%s",
litellm_call_id,
extra=MappingProxyType({"litellm_call_id": litellm_call_id}),
)
return
log_fn: Final = (
@ -1462,7 +1466,12 @@ def _log_llm_api_exception(e: Exception) -> None:
if is_expected_client_error(e) and not litellm.log_client_error_tracebacks
else verbose_proxy_logger.exception
)
log_fn("litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - %s", e)
log_fn(
"litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - litellm_call_id=%s - %s",
litellm_call_id,
e,
extra=MappingProxyType({"litellm_call_id": litellm_call_id}),
)
async def _cancel_llm_call_on_client_disconnect(
@ -2338,9 +2347,8 @@ class ProxyBaseLLMRequestProcessing:
"""
Common request processing logic for both chat completions and responses API endpoints
"""
requested_model_from_client: Final[str | None] = (
self.data.get("model") if isinstance(self.data.get("model"), str) else None
)
client_model: Final = get_client_requested_model(request) or self.data.get("model")
requested_model_from_client: Final[str | None] = client_model if isinstance(client_model, str) else None
self._debug_log_request_payload()
if skip_pre_call_logic:
@ -3421,7 +3429,11 @@ class ProxyBaseLLMRequestProcessing:
version: str | None = None,
):
"""Raises ProxyException (OpenAI API compatible) if an exception is raised"""
_log_llm_api_exception(e)
logging_obj: Final[LiteLLMLoggingObj | None] = self.data.get("litellm_logging_obj", None)
_log_llm_api_exception(
e,
(logging_obj.litellm_call_id if logging_obj is not None else None) or self.data.get("litellm_call_id"),
)
# Allow callbacks to transform the error response
transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,

View file

@ -9,7 +9,7 @@ from fastapi import Request, UploadFile, status
from typing_extensions import NotRequired, ReadOnly, Required
from litellm._logging import verbose_proxy_logger
from litellm.constants import MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB
from litellm.constants import CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB
from litellm.proxy._types import ProxyException
from litellm.proxy.common_utils.callback_utils import (
get_metadata_variable_name_from_kwargs,
@ -235,6 +235,13 @@ def _safe_get_request_parsed_body(request: Request | None) -> dict | None:
return None
def get_client_requested_model(request: Request | None) -> str | None:
if request is None or not hasattr(request, "scope"):
return None
model: Final = request.scope.get(CLIENT_REQUESTED_MODEL_SCOPE_KEY)
return model if isinstance(model, str) else None
def _safe_get_request_query_params(request: Request | None) -> dict:
if request is None:
return {}
@ -259,6 +266,24 @@ def _safe_set_request_parsed_body(
verbose_proxy_logger.debug("Unexpected error setting request parsed body - %s", e)
def rewrite_request_model(
request_data: dict[str, object], # mutable-ok: the request body is rewritten in place for every downstream reader
request: Request | None,
model: str,
) -> None:
"""Point the auth-time payload, the parsed-body cache, ``request.json()`` and ``request.body()`` at ``model``.
The cache and raw body keep only the keys the client sent, not params auth merged into ``request_data``.
"""
request_data["model"] = model
if request is None:
return
cached_body: Final = _safe_get_request_parsed_body(request=request)
body: Final = {**cached_body, "model": model} if cached_body is not None else request_data
_safe_set_request_parsed_body(request=request, parsed_body=body)
request._json = body
request._body = orjson.dumps(body)
def _safe_get_request_headers(request: Request | None) -> dict:
"""
[Non-Blocking] Safely get the request headers.

View file

@ -16,6 +16,7 @@ from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload
from urllib.parse import quote, unquote
import litellm
from litellm._logging import verbose_proxy_logger
@ -85,6 +86,10 @@ else:
RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value})
def _org_member_transaction_key(org_id: str, user_id: str) -> str:
return f"organization_id::{quote(org_id, safe='')}::user_id::{quote(user_id, safe='')}"
def _is_batch_cost_row(payload: SpendLogsPayload) -> bool:
return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success"
@ -110,6 +115,7 @@ class _SpendBatch(Protocol):
litellm_teamtable: BatchTable
litellm_teammembership: BatchTable
litellm_organizationtable: BatchTable
litellm_organizationmembership: BatchTable
litellm_tagtable: BatchTable
litellm_agentstable: BatchTable
litellm_modelaccessgroupbudgettable: BatchTable
@ -666,6 +672,7 @@ class DBSpendUpdateWriter:
await self._update_org_db(
response_cost=response_cost,
org_id=org_id,
user_id=user_id,
prisma_client=prisma_client,
)
except Exception:
@ -900,6 +907,7 @@ class DBSpendUpdateWriter:
self,
response_cost: float | None,
org_id: str | None,
user_id: str | None,
prisma_client: PrismaClient | None,
):
try:
@ -916,6 +924,15 @@ class DBSpendUpdateWriter:
response_cost=response_cost,
)
)
if user_id is not None:
await self.spend_update_queue.add_update(
update=SpendUpdateQueueItem(
entity_type=Litellm_EntityType.ORGANIZATION_MEMBER,
entity_id=_org_member_transaction_key(org_id, user_id),
response_cost=response_cost,
)
)
except Exception as e:
spend_log_error(
"Spend tracking - failed to enqueue org spend update. org_id=%s, response_cost=%s - %s",
@ -1163,14 +1180,15 @@ class DBSpendUpdateWriter:
if db_spend_update_transactions is not None:
verbose_proxy_logger.info(
"Spend tracking - committing spend updates from Redis to DB: "
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d, agents=%d, "
"model_access_groups=%d",
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, tags=%d, "
"agents=%d, model_access_groups=%d",
len(db_spend_update_transactions.get("key_list_transactions") or {}),
len(db_spend_update_transactions.get("user_list_transactions") or {}),
len(db_spend_update_transactions.get("team_list_transactions") or {}),
len(db_spend_update_transactions.get("org_list_transactions") or {}),
len(db_spend_update_transactions.get("end_user_list_transactions") or {}),
len(db_spend_update_transactions.get("team_member_list_transactions") or {}),
len(db_spend_update_transactions.get("org_member_list_transactions") or {}),
len(db_spend_update_transactions.get("tag_list_transactions") or {}),
len(db_spend_update_transactions.get("agent_list_transactions") or {}),
len(db_spend_update_transactions.get("model_access_group_list_transactions") or {}),
@ -1708,6 +1726,29 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
)
org_member_list_transactions: Final = db_spend_update_transactions.get("org_member_list_transactions")
verbose_proxy_logger.debug("Org Membership Spend transactions: %s", org_member_list_transactions)
if org_member_list_transactions is not None and len(org_member_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(prisma_client) as transaction, transaction.batch_() as batcher:
for key, response_cost in sorted(org_member_list_transactions.items()):
_, quoted_org_id, _, quoted_user_id = key.split("::")
batcher.litellm_organizationmembership.update_many(
where={"organization_id": unquote(quoted_org_id), "user_id": unquote(quoted_user_id)},
data={"spend": {"increment": response_cost}},
)
break
except Exception as e:
await self._handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE TAG TABLE ###
tag_list_transactions: Final = db_spend_update_transactions["tag_list_transactions"]
await DBSpendUpdateWriter._update_entity_spend_in_db(

View file

@ -69,6 +69,7 @@ _SpendTransactionField: TypeAlias = Literal[
"team_list_transactions",
"team_member_list_transactions",
"org_list_transactions",
"org_member_list_transactions",
"tag_list_transactions",
"agent_list_transactions",
"model_access_group_list_transactions",
@ -81,6 +82,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = (
"team_list_transactions",
"team_member_list_transactions",
"org_list_transactions",
"org_member_list_transactions",
"tag_list_transactions",
"agent_list_transactions",
"model_access_group_list_transactions",
@ -412,6 +414,10 @@ class RedisUpdateBuffer:
Litellm_EntityType.ORGANIZATION,
db_spend_update_transactions.get("org_list_transactions"),
),
(
Litellm_EntityType.ORGANIZATION_MEMBER,
db_spend_update_transactions.get("org_member_list_transactions"),
),
(
Litellm_EntityType.TAG,
db_spend_update_transactions.get("tag_list_transactions"),
@ -876,6 +882,9 @@ class RedisUpdateBuffer:
list_of_transactions, "team_member_list_transactions"
),
org_list_transactions=_merged_entity_transactions(list_of_transactions, "org_list_transactions"),
org_member_list_transactions=_merged_entity_transactions(
list_of_transactions, "org_member_list_transactions"
),
tag_list_transactions=_merged_entity_transactions(list_of_transactions, "tag_list_transactions"),
agent_list_transactions=_merged_entity_transactions(list_of_transactions, "agent_list_transactions"),
model_access_group_list_transactions=_merged_entity_transactions(

View file

@ -137,6 +137,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
team_list_transactions={},
team_member_list_transactions={},
org_list_transactions={},
org_member_list_transactions={},
tag_list_transactions={},
agent_list_transactions={},
model_access_group_list_transactions={},
@ -150,6 +151,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
Litellm_EntityType.TEAM: "team_list_transactions",
Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions",
Litellm_EntityType.ORGANIZATION: "org_list_transactions",
Litellm_EntityType.ORGANIZATION_MEMBER: "org_member_list_transactions",
Litellm_EntityType.TAG: "tag_list_transactions",
Litellm_EntityType.AGENT: "agent_list_transactions",
Litellm_EntityType.MODEL_ACCESS_GROUP: "model_access_group_list_transactions",
@ -188,6 +190,8 @@ class SpendUpdateQueue(BaseUpdateQueue):
transactions_dict = db_spend_update_transactions["team_member_list_transactions"]
elif dict_key == "org_list_transactions":
transactions_dict = db_spend_update_transactions["org_list_transactions"]
elif dict_key == "org_member_list_transactions":
transactions_dict = db_spend_update_transactions["org_member_list_transactions"]
elif dict_key == "tag_list_transactions":
transactions_dict = db_spend_update_transactions["tag_list_transactions"]
elif dict_key == "agent_list_transactions":

View file

@ -41,6 +41,7 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import (
ESTIMATED_OUTPUT_TOKENS_FIELD,
get_estimated_output_tokens,
get_key_own_model_rate_limit,
get_key_tag_rpm_limit,
get_model_rate_limit_from_metadata,
)
@ -2892,41 +2893,67 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return batch_limiter
return None
def _key_owns_model_limit(
self,
user_api_key_dict: UserAPIKeyAuth,
requested_model: str,
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
) -> bool:
key_own_limits: Final = get_key_own_model_rate_limit(user_api_key_dict, rate_limit_key)
return key_own_limits is not None and key_own_limits.get(requested_model) is not None
def _inherited_team_model_limit(
self,
user_api_key_dict: UserAPIKeyAuth,
requested_model: str,
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
) -> int | None:
team_limits: Final = get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", rate_limit_key)
team_limit: Final = team_limits.get(requested_model) if team_limits else None
if team_limit is None:
return None
if self._key_owns_model_limit(user_api_key_dict, requested_model, rate_limit_key):
return None
return team_limit
def _key_owns_model_tpm_limit_from_request_metadata(
self,
request_metadata: Mapping[str, object],
model_group: str | None,
) -> bool:
if model_group is None:
return False
key_view: Final = UserAPIKeyAuth.model_validate(
{
"metadata": request_metadata.get("user_api_key_metadata") or {},
"model_max_budget": request_metadata.get("user_api_key_model_max_budget") or {},
}
)
return self._key_owns_model_limit(key_view, model_group, "model_tpm_limit")
def _add_team_model_rate_limit_descriptor_from_metadata(
self,
user_api_key_dict: UserAPIKeyAuth,
requested_model: str | None,
descriptors: list[RateLimitDescriptor],
) -> None:
"""Add team model rate limit descriptor from team_metadata if applicable."""
if (
get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_rpm_limit") is not None
or get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_tpm_limit") is not None
):
_tpm_limit_for_team_model: Final = (
get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_tpm_limit") or {}
if requested_model is None:
return
team_rpm_limit: Final = self._inherited_team_model_limit(user_api_key_dict, requested_model, "model_rpm_limit")
team_tpm_limit: Final = self._inherited_team_model_limit(user_api_key_dict, requested_model, "model_tpm_limit")
if team_rpm_limit is None and team_tpm_limit is None:
return
descriptors.append(
RateLimitDescriptor(
key="model_per_team",
value=f"{user_api_key_dict.team_id}:{requested_model}",
rate_limit={
"requests_per_unit": team_rpm_limit,
"tokens_per_unit": team_tpm_limit,
"window_size": self.window_size,
},
)
_rpm_limit_for_team_model: Final = (
get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_rpm_limit") or {}
)
should_check_rate_limit: Final = (
requested_model in _tpm_limit_for_team_model or requested_model in _rpm_limit_for_team_model
)
if should_check_rate_limit and requested_model is not None:
model_specific_tpm_limit: Final = _tpm_limit_for_team_model.get(requested_model)
model_specific_rpm_limit: Final = _rpm_limit_for_team_model.get(requested_model)
descriptors.append(
RateLimitDescriptor(
key="model_per_team",
value=f"{user_api_key_dict.team_id}:{requested_model}",
rate_limit={
"requests_per_unit": model_specific_rpm_limit,
"tokens_per_unit": model_specific_tpm_limit,
"window_size": self.window_size,
},
)
)
)
def _add_project_model_rate_limit_descriptor_from_metadata(
self,
@ -4459,6 +4486,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
kwargs=kwargs,
model_group=reconcile_model,
)
charged_targets: Final = (
[target for target in targets if target[0] != "model_per_team"]
if self._key_owns_model_tpm_limit_from_request_metadata(request_metadata, reconcile_model)
else targets
)
if reserved_tokens > 0 and total_tokens < reserved_tokens:
verbose_proxy_logger.debug(
"Releasing unused TPM budget on success: reserved=%s, actual=%s, release=%s",
@ -4468,7 +4500,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
pipeline_operations.extend(
self._build_reservation_aware_tpm_ops(
targets=targets,
targets=charged_targets,
reserved_scopes=reserved_scopes,
actual_tokens=total_tokens,
reserved_tokens=reserved_tokens,

View file

@ -28,6 +28,9 @@ class _JWTKeyMappingRecord(Protocol):
@property
def id(self) -> str: ...
@property
def jwt_issuer(self) -> str: ...
@property
def jwt_claim_name(self) -> str: ...
@ -78,6 +81,7 @@ def _to_response(mapping: _JWTKeyMappingRecord) -> JWTKeyMappingResponse:
"""Convert a Prisma mapping object to a safe response (no hashed token)."""
return JWTKeyMappingResponse(
id=mapping.id,
jwt_issuer=mapping.jwt_issuer or None,
jwt_claim_name=mapping.jwt_claim_name,
jwt_claim_value=mapping.jwt_claim_value,
description=mapping.description,
@ -109,6 +113,7 @@ async def create_jwt_key_mapping(
try:
hashed_key: Final = hash_token(data.key)
create_data: Final = {
"jwt_issuer": data.jwt_issuer or "",
"jwt_claim_name": data.jwt_claim_name,
"jwt_claim_value": data.jwt_claim_value,
"token": hashed_key,
@ -120,7 +125,7 @@ async def create_jwt_key_mapping(
new_mapping: Final = await _mapping_table(prisma_client).create(data=create_data)
cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value)
cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value, data.jwt_issuer)
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
return _to_response(new_mapping)
@ -131,7 +136,10 @@ async def create_jwt_key_mapping(
if "unique" in error_str or "p2002" in error_str:
raise HTTPException(
status_code=409,
detail=f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' already exists.",
detail=(
f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' "
f"already exists for issuer '{data.jwt_issuer}'."
),
)
if "foreign" in error_str or "p2003" in error_str:
raise HTTPException(
@ -161,6 +169,9 @@ async def update_jwt_key_mapping(
update_data: Final = data.model_dump(exclude_unset=True, exclude={"id", "key"})
if data.key is not None:
update_data["token"] = hash_token(data.key)
if "jwt_issuer" in update_data:
# DB column is NOT NULL (see schema.prisma); "" is the global/unscoped sentinel.
update_data["jwt_issuer"] = update_data["jwt_issuer"] or ""
update_data["updated_by"] = user_api_key_dict.user_id
try:
@ -178,9 +189,11 @@ async def update_jwt_key_mapping(
# Evict only after the write commits: a concurrent request between an
# early eviction and the commit would re-cache the old mapping and keep
# it authorized until TTL.
old_cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value)
old_cache_key: Final = jwt_key_mapping_cache_key(
old_mapping.jwt_claim_name, old_mapping.jwt_claim_value, old_mapping.jwt_issuer
)
new_cache_key: Final = jwt_key_mapping_cache_key(
updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value
updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value, updated_mapping.jwt_issuer
)
cache_keys: Final = (old_cache_key,) if old_cache_key == new_cache_key else (old_cache_key, new_cache_key)
await evict_and_broadcast(cache_keys=cache_keys, user_api_key_cache=user_api_key_cache)
@ -227,7 +240,9 @@ async def delete_jwt_key_mapping(
# Evict only after the row is gone, else a concurrent request can
# re-cache the deleted mapping and keep it authorized until TTL.
cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value)
cache_key: Final = jwt_key_mapping_cache_key(
old_mapping.jwt_claim_name, old_mapping.jwt_claim_value, old_mapping.jwt_issuer
)
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
return {"status": "success"}
except HTTPException:

View file

@ -36,6 +36,7 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
from litellm.llms.azure.passthrough.transformation import foreign_azure_deployment
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
from litellm.proxy._types import *
@ -77,6 +78,7 @@ from litellm.proxy.vector_store_endpoints.utils import (
from litellm.secret_managers.main import get_secret_str, str_to_bool
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
)
from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials
@ -1322,7 +1324,7 @@ def _resolve_vertex_model_from_router(
endpoint: str,
vertex_project: str | None,
vertex_location: str | None,
) -> tuple[str, str, str | None, str | None]:
) -> tuple[str, str, str | None, str | None, Mapping[str, object] | None]:
"""
Resolve Vertex AI model configuration from router.
@ -1335,18 +1337,21 @@ def _resolve_vertex_model_from_router(
vertex_location: Current vertex location (may be from URL)
Returns:
tuple of (encoded_endpoint, endpoint, vertex_project, vertex_location)
with resolved values from router config
tuple of (encoded_endpoint, endpoint, vertex_project, vertex_location, deployment_model_info)
with resolved values from router config; deployment_model_info is the resolved
deployment's `model_info`, or None when no deployment matched
"""
if not llm_router:
return encoded_endpoint, endpoint, vertex_project, vertex_location
return encoded_endpoint, endpoint, vertex_project, vertex_location, None
try:
deployment: Final = llm_router.get_available_deployment_for_pass_through(model=model_id)
if not deployment:
return encoded_endpoint, endpoint, vertex_project, vertex_location
return encoded_endpoint, endpoint, vertex_project, vertex_location, None
litellm_params: Final = deployment.get("litellm_params", {})
model_info: Final = deployment.get("model_info")
deployment_model_info: Final = model_info if isinstance(model_info, Mapping) else None
# Always override with router config values (they take precedence over URL values)
config_vertex_project: Final = litellm_params.get("vertex_project")
@ -1387,10 +1392,11 @@ def _resolve_vertex_model_from_router(
encoded_endpoint = encoded_endpoint.replace(model_id, actual_model)
endpoint = endpoint.replace(model_id, actual_model)
return encoded_endpoint, endpoint, vertex_project, vertex_location, deployment_model_info
except Exception as e:
verbose_proxy_logger.debug("Error resolving vertex model from router for model %s: %s", model_id, e)
return encoded_endpoint, endpoint, vertex_project, vertex_location
return encoded_endpoint, endpoint, vertex_project, vertex_location, None
def _is_bedrock_agent_runtime_route(endpoint: str) -> bool:
@ -1545,6 +1551,26 @@ async def _relay_azure_router_model(
"put the model group name in the deployments segment"
}
raise HTTPException(status_code=400, detail=rejection)
return await _relay_router_model(
llm_router=llm_router,
model=model,
endpoint=endpoint,
request=request,
request_body=request_body,
is_streaming_request=is_streaming_request,
user_api_key_dict=user_api_key_dict,
)
async def _relay_router_model(
llm_router: litellm.Router,
model: str,
endpoint: str,
request: Request,
request_body: Mapping[str, object],
is_streaming_request: bool,
user_api_key_dict: UserAPIKeyAuth,
) -> Response:
try:
result: Final = await llm_router.allm_passthrough_route(
model=model,
@ -1594,6 +1620,65 @@ async def _relay_azure_router_model(
)
@router.api_route(
"/nvidia_nim/{endpoint:path}",
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
tags=["NVIDIA NIM Pass-through", "pass-through"],
)
async def nvidia_nim_proxy_route(
endpoint: str,
request: Request,
fastapi_response: Response,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Relay a native NVIDIA NIM request through a LiteLLM model group.
`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's
`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through
virtual key auth, model access checks, and spend logging.
"""
from litellm.proxy.proxy_server import llm_router
return await relay_nvidia_nim_request(
llm_router=llm_router,
endpoint=endpoint,
request=request,
request_body=await get_request_body(request),
user_api_key_dict=user_api_key_dict,
)
async def relay_nvidia_nim_request(
llm_router: litellm.Router | None,
endpoint: str,
request: Request,
request_body: Mapping[str, object],
user_api_key_dict: UserAPIKeyAuth,
) -> Response:
model_group: Final = nvidia_nim_model_group_in_path(endpoint, llm_router.get_model_list()) if llm_router else None
if llm_router is None or model_group is None:
rejection: Final[RelayRejection] = {
"error": "no NVIDIA NIM model group in the path; call /nvidia_nim/{model_group}/v1/infer with a model "
"group from your `model_list` whose deployments all use `nvidia_nim/` models"
}
raise HTTPException(status_code=400, detail=rejection)
is_streaming_request: Final = is_passthrough_request_streaming(request_body)
return await open_sse_before_first_byte(
_relay_router_model(
llm_router=llm_router,
model=model_group,
endpoint=endpoint,
request=request,
request_body=request_body,
is_streaming_request=is_streaming_request,
user_api_key_dict=user_api_key_dict,
),
ping_interval_seconds=(litellm.sse_keepalive_ping_interval_seconds if is_streaming_request else None),
)
@router.api_route(
"/azure_ai/{endpoint:path}",
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
@ -2134,6 +2219,7 @@ async def _base_vertex_proxy_route(
endpoint,
vertex_project,
vertex_location,
deployment_model_info,
) = _resolve_vertex_model_from_router(
model_id=model_id,
llm_router=llm_router,
@ -2142,6 +2228,8 @@ async def _base_vertex_proxy_route(
vertex_project=vertex_project,
vertex_location=vertex_location,
)
if deployment_model_info:
setattr(request.state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, deployment_model_info)
vertex_credentials: Final = passthrough_endpoint_router.get_vertex_credentials(
project_id=vertex_project,

View file

@ -12,6 +12,9 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator as GeminiModelResponseIterator,
)
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import (
VertexPassthroughLoggingHandler,
)
from litellm.types.utils import (
ModelResponse,
TextCompletionResponse,
@ -40,6 +43,17 @@ class GeminiPassthroughLoggingHandler:
request_body: dict,
**kwargs,
) -> PassThroughEndpointLoggingTypedDict:
if VertexPassthroughLoggingHandler.is_interactions_route(url_route):
return VertexPassthroughLoggingHandler.interactions_passthrough_handler(
httpx_response=httpx_response,
request_body=request_body,
logging_obj=logging_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
custom_llm_provider="gemini",
vertex_location=None,
)
if "predictLongRunning" in url_route:
model = GeminiPassthroughLoggingHandler.extract_model_from_url(url_route)

View file

@ -1,15 +1,20 @@
import asyncio
import re
from collections.abc import Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, cast
from typing import TYPE_CHECKING, Any, Final, Literal, cast
from urllib.parse import urlparse
import httpx
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import VERTEX_BATCH_PREDICTION_JOBS_ROUTE
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import (
InteractionsUsageObjectTransformation,
)
from litellm.llms.vertex_ai.common_utils import (
get_vertex_ai_lyria_generation_cost,
get_vertex_location_from_url,
@ -49,8 +54,73 @@ else:
EndpointType = Any
_VERTEX_INTERACTIONS_PATH: Final = re.compile(r"/projects/[^/]+/locations/[^/]+/interactions/?$")
_INTERACTIONS_RESPONSE_BODY: Final = TypeAdapter(dict[str, object])
def _interactions_model(
response_body: Mapping[str, object],
request_body: Mapping[str, object] | None,
) -> str | None:
response_model: Final = response_body.get("model")
if isinstance(response_model, str) and response_model:
return response_model
request_model: Final = (request_body or {}).get("model")
if isinstance(request_model, str) and request_model:
return request_model
return None
class VertexPassthroughLoggingHandler:
@staticmethod
def is_interactions_route(url_route: str) -> bool:
return urlparse(url_route).path.rstrip("/").endswith("/interactions")
@staticmethod
def is_vertex_interactions_route(url_route: str) -> bool:
return _VERTEX_INTERACTIONS_PATH.search(urlparse(url_route).path) is not None
@staticmethod
def interactions_passthrough_handler(
httpx_response: httpx.Response,
request_body: Mapping[str, object] | None,
logging_obj: LiteLLMLoggingObj,
kwargs: dict[str, object],
start_time: datetime,
end_time: datetime,
custom_llm_provider: Literal["vertex_ai", "gemini"],
vertex_location: str | None,
) -> PassThroughEndpointLoggingTypedDict:
response_body: Final = _INTERACTIONS_RESPONSE_BODY.validate_python(httpx_response.json())
usage_object: Final = response_body.get("usage")
model: Final = _interactions_model(response_body, request_body)
if model is None or not InteractionsUsageObjectTransformation.is_interactions_usage_object(usage_object):
return {"result": None, "kwargs": kwargs}
litellm_model_response: Final = ModelResponse(
model=model,
usage=InteractionsUsageObjectTransformation.transform_interactions_usage_object(
cast(Mapping[str, Any], usage_object)
),
)
logging_obj.custom_llm_provider = custom_llm_provider
logging_kwargs: Final = (
VertexPassthroughLoggingHandler._create_vertex_response_logging_payload_for_generate_content(
litellm_model_response=litellm_model_response,
model=model,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
vertex_location=vertex_location,
)
)
return {
"result": litellm_model_response,
"kwargs": {**logging_kwargs, "custom_llm_provider": custom_llm_provider},
}
@staticmethod
def vertex_passthrough_handler(
httpx_response: httpx.Response,
@ -66,6 +136,17 @@ class VertexPassthroughLoggingHandler:
vertex_location: Final = get_vertex_location_from_url(url_route)
if vertex_location is not None:
logging_obj.optional_params["vertex_location"] = vertex_location
if VertexPassthroughLoggingHandler.is_interactions_route(url_route):
return VertexPassthroughLoggingHandler.interactions_passthrough_handler(
httpx_response=httpx_response,
request_body=request_body,
logging_obj=logging_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
custom_llm_provider="vertex_ai",
vertex_location=vertex_location,
)
if "predictLongRunning" in url_route:
model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route)

View file

@ -96,6 +96,7 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
EndpointType,
@ -613,6 +614,12 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
_metadata.update(
LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict)
)
_request_state: Final = getattr(request, "state", None)
deployment_model_info: Final = getattr(
_request_state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, None
)
if isinstance(deployment_model_info, Mapping):
_metadata["model_info"] = dict(deployment_model_info)
kwargs: Final = {
"litellm_params": {
@ -2002,6 +2009,8 @@ def create_pass_through_route(
delattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY)
if hasattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY):
delattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY)
if hasattr(request.state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY):
delattr(request.state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY)
# The upstream withholds its response headers until its first token, so
# the whole time-to-first-token is spent inside _relay with nothing on

View file

@ -361,7 +361,9 @@ class PassThroughEndpointLogging:
def is_vertex_route(self, url_route: str) -> bool:
if any(f":{method}" in url_route for method in self.TRACKED_VERTEX_METHOD_ROUTES):
return True
return any(resource in url_route for resource in self.TRACKED_VERTEX_RESOURCE_ROUTES)
if any(resource in url_route for resource in self.TRACKED_VERTEX_RESOURCE_ROUTES):
return True
return VertexPassthroughLoggingHandler.is_vertex_interactions_route(url_route)
def is_anthropic_route(self, url_route: str):
for route in self.TRACKED_ANTHROPIC_ROUTES:
@ -434,8 +436,12 @@ class PassThroughEndpointLogging:
def is_gemini_route(self, url_route: str, custom_llm_provider: str | None = None):
"""Check if the URL route is a Gemini API route."""
if custom_llm_provider != "gemini":
return False
if VertexPassthroughLoggingHandler.is_interactions_route(url_route):
return True
for route in self.TRACKED_GEMINI_ROUTES:
if route in url_route and custom_llm_provider == "gemini":
if route in url_route:
return True
return False

View file

@ -261,7 +261,7 @@ class ProxyInitializationHelpers:
import uvicorn
import litellm
from litellm._logging import _get_uvicorn_json_log_config
from litellm._logging import _get_uvicorn_json_log_config, resolve_log_level
uvicorn_args: Final = {
"app": "litellm.proxy.proxy_server:app",
@ -275,6 +275,8 @@ class ProxyInitializationHelpers:
elif litellm.json_logs:
# Use JSON log config for uvicorn to ensure all logs (including exceptions) are JSON
uvicorn_args["log_config"] = _get_uvicorn_json_log_config()
elif litellm_log := os.environ.get("LITELLM_LOG"):
uvicorn_args["log_level"] = resolve_log_level(litellm_log)
if keepalive_timeout is not None:
uvicorn_args["timeout_keep_alive"] = keepalive_timeout
if timeout_worker_healthcheck is not None:

View file

@ -7080,8 +7080,19 @@ class ProxyConfig:
## PASS-THROUGH ENDPOINTS ##
if "pass_through_endpoints" in _general_settings:
general_settings["pass_through_endpoints"] = _general_settings["pass_through_endpoints"]
await initialize_pass_through_endpoints(pass_through_endpoints=general_settings["pass_through_endpoints"])
db_pass_through_endpoints: Final = _general_settings["pass_through_endpoints"]
db_pass_through_paths: Final = frozenset(
endpoint.get("path") for endpoint in db_pass_through_endpoints if isinstance(endpoint, dict)
)
general_settings["pass_through_endpoints"] = [
*db_pass_through_endpoints,
*(
endpoint
for endpoint in config_passthrough_endpoints or ()
if endpoint.get("path") not in db_pass_through_paths
),
]
await initialize_pass_through_endpoints(pass_through_endpoints=db_pass_through_endpoints)
## UI ACCESS MODE ##
if "ui_access_mode" in _general_settings:
@ -17465,6 +17476,16 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: Final[dict[str, GeneralSettingsUILiteLLMFie
"tab": "prompt_caching",
"description": "Empty uses Anthropic's 5m default. 1h suits long sessions but doubles the cache write cost.",
},
"openai_system_messages_first": {
"type": "Boolean",
"tab": "prompt_caching",
"description": (
"Moves system and developer messages to the front of the messages array on OpenAI and "
"Azure OpenAI chat completions requests, keeping their relative order. OpenAI's prompt cache "
"matches on the exact prefix, so a system message that arrives mid-conversation otherwise "
"breaks the cached prefix on every turn."
),
},
"budget_rollover": { # mutable-ok: registry literal, frozen with its siblings below
"type": "Boolean",
"description": (

View file

@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
model LiteLLM_JWTKeyMapping {
id String @id @default(uuid())
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
// Not nullable: Postgres unique constraints treat every NULL as
// distinct, so a nullable column would let multiple unscoped
// mappings collide on the same claim without a constraint violation.
jwt_claim_name String // e.g. "sub", "email"
jwt_claim_value String // The claim value to match
token String // Hashed virtual key (FK)
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
}
// Deprecated keys during grace period - allows old key to work until revoke_at

View file

@ -176,6 +176,7 @@ class _SessionSpendRow(TypedDict):
api_key: ReadOnly[str]
session_total_count: ReadOnly[int]
session_total_spend: float
session_total_duration_ms: ReadOnly[int]
mcp_tool_call_count: int
mcp_tool_call_spend: float
session_cache_hit_count: ReadOnly[int]
@ -194,6 +195,7 @@ _SESSION_MODEL_NAME_MAX_LEN: Final = 256
class _SessionSpendStats(NamedTuple):
session_total_count: int
session_total_spend: float
session_total_duration_ms: int
mcp_tool_call_count: int
mcp_tool_call_spend: float
session_cache_hit_count: int
@ -4543,6 +4545,12 @@ async def _build_ui_spend_logs_response(
SELECT session_id, api_key,
COUNT(*)::int AS session_total_count,
COALESCE(SUM(spend), 0)::double precision AS session_total_spend,
COALESCE(SUM(
COALESCE(
request_duration_ms,
(EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER
)
), 0)::bigint AS session_total_duration_ms,
COUNT(*) FILTER (
WHERE call_type IN {_MCP_CALL_TYPES_SQL}
)::int AS mcp_tool_call_count,
@ -4584,6 +4592,7 @@ async def _build_ui_spend_logs_response(
(row["session_id"], row["api_key"]): _SessionSpendStats(
session_total_count=int(row.get("session_total_count") or 0),
session_total_spend=float(row.get("session_total_spend") or 0.0),
session_total_duration_ms=int(row.get("session_total_duration_ms") or 0),
mcp_tool_call_count=int(row.get("mcp_tool_call_count") or 0),
mcp_tool_call_spend=float(row.get("mcp_tool_call_spend") or 0.0),
session_cache_hit_count=int(row.get("session_cache_hit_count") or 0),
@ -4615,6 +4624,7 @@ async def _build_ui_spend_logs_response(
row_dict["session_total_count"] = session_stats.session_total_count if session_stats else 1
if session_stats:
row_dict["session_total_spend"] = session_stats.session_total_spend
row_dict["session_total_duration_ms"] = session_stats.session_total_duration_ms
if session_stats.mcp_tool_call_count:
row_dict["mcp_tool_call_count"] = session_stats.mcp_tool_call_count
row_dict["mcp_tool_call_spend"] = session_stats.mcp_tool_call_spend

View file

@ -102,6 +102,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import (
mask_sensitive_structure,
)
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
from litellm.llms.base_llm.passthrough.transformation import replace_path_segment
from litellm.llms.base_llm.vector_store.transformation import (
RouterVectorStoreEmbeddingExecutor,
@ -10818,6 +10819,7 @@ class Router:
"model_group": user_facing_model_group_name,
"providers": [llm_provider],
**model_info,
"supports_fast_mode": True,
"supported_reasoning_efforts": None,
}
)
@ -10896,6 +10898,9 @@ class Router:
if model_info.get("rpm", None) is not None and _deployment_rpm is None:
_deployment_rpm = model_info.get("rpm")
model_group_info.supports_fast_mode = model_group_info.supports_fast_mode and (
AnthropicModelInfo.supports_fast_mode(litellm_model, llm_provider)
)
deployment_reasoning_efforts = (
resolve_supported_reasoning_efforts( # rebind-ok: recalculated per deployment
model_info, deployment_is_mapped=deployment_is_mapped

View file

@ -11,6 +11,9 @@ LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY: Final = "litellm_pass_through_custom
# exact byte/string body, such as AWS SigV4-signed requests.
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY: Final = "litellm_pass_through_raw_body"
# `model_info` of the router deployment a provider route (e.g. Vertex) resolved for this request.
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY: Final = "litellm_pass_through_deployment_model_info"
# Attribute set on the FastAPI endpoint function of every user-defined pass-through
# route. Auth reads it off the dispatched endpoint (``request.scope["endpoint"]``) to
# decide whether a request body ``model`` names an upstream model rather than a

View file

@ -722,6 +722,7 @@ class ModelGroupInfo(BaseModel):
supports_url_context: bool = Field(default=False)
supports_reasoning: bool = Field(default=False)
supports_function_calling: bool = Field(default=False)
supports_fast_mode: bool = Field(default=False)
supported_reasoning_efforts: tuple[str, ...] | None = Field(default=None)
supported_openai_params: list[str] | None = Field(default=[])
configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None

View file

@ -4275,6 +4275,7 @@ class LiteLLMRealtimeStreamLoggingObject(LiteLLMPydanticObjectBase):
# rate_limits.updated), blocks the event loop, and discards the session usage.
results: SkipValidation[OpenAIRealtimeStreamList]
usage: Usage
service_tier: str | None = None
_hidden_params: dict = {}
@field_serializer("results")

View file

@ -8998,6 +8998,12 @@ class ProviderConfigManager:
)
return WatsonxPassthroughConfig()
elif LlmProviders.NVIDIA_NIM == provider:
from litellm.llms.nvidia_nim.passthrough.transformation import (
NvidiaNimPassthroughConfig,
)
return NvidiaNimPassthroughConfig()
return None
@staticmethod

File diff suppressed because it is too large Load diff

View file

@ -716,6 +716,9 @@
"supports_embedding_image_input": {
"type": "boolean"
},
"supports_fast_mode": {
"type": "boolean"
},
"supports_forced_tool_use": {
"type": "boolean"
},

View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.102.0"
version = "1.103.0"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.15"
@ -67,8 +67,8 @@ proxy = [
"azure-identity>=1.25.2,<2.0",
"azure-storage-blob>=12.28.0,<13.0",
"mcp>=1.28.1,<2.0",
"litellm-proxy-extras==0.4.97",
"litellm-enterprise==0.1.67",
"litellm-proxy-extras==0.4.98",
"litellm-enterprise==0.1.68",
"RestrictedPython>=8.5,<9.0",
"rich>=13.9.4,<14.0",
"InquirerPy>=0.3.4,<1.0",
@ -290,6 +290,7 @@ editable-profile = "dev"
include = [
"litellm/proxy/_experimental/out/**",
"litellm/router_strategy/complexity_router/artifacts/*.json",
"litellm/proxy/client/cli/commands/codex_base_instructions.md",
]
exclude = [
"litellm/proxy/enterprise",
@ -331,7 +332,7 @@ members = ["enterprise", "litellm-proxy-extras"]
profile = "black"
[tool.commitizen]
version = "1.102.0"
version = "1.103.0"
version_files = [
"pyproject.toml:^version",
]

View file

@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
model LiteLLM_JWTKeyMapping {
id String @id @default(uuid())
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
// Not nullable: Postgres unique constraints treat every NULL as
// distinct, so a nullable column would let multiple unscoped
// mappings collide on the same claim without a constraint violation.
jwt_claim_name String // e.g. "sub", "email"
jwt_claim_value String // The claim value to match
token String // Hashed virtual key (FK)
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
}
// Deprecated keys during grace period - allows old key to work until revoke_at

View file

@ -67,6 +67,7 @@ IGNORE_FUNCTIONS = [
"_redact_agent_params_tree", # max depth set (default 10), same shape as _redact_sensitive_litellm_params.
"_restore_redacted_nested_value", # max depth set (default 10), mirrors _redact_agent_params_tree on the write side.
"_unqualified", # bounded by the qualifier depth of a static TypedDict annotation (Annotated, Required/NotRequired, ReadOnly around one type, no cycles possible).
"completion_cost", # max depth 1: recursion only fires for mixed-tier Responses WS logging objects, and each split part carries a single service_tier so _split_responses_ws_logging_object_by_service_tier returns None.
]

View file

@ -0,0 +1,37 @@
import importlib
import logging
from collections.abc import Iterator
import pytest
import litellm_proxy_extras._logging as extras_logging
@pytest.fixture
def fresh_extras_logger() -> Iterator[logging.Logger]:
logger = logging.getLogger("litellm_proxy_extras")
saved_handlers = logger.handlers[:]
saved_level = logger.level
logger.handlers[:] = []
try:
yield logger
finally:
logger.handlers[:] = saved_handlers
logger.setLevel(saved_level)
def test_litellm_log_error_silences_extras_info_lines(monkeypatch, fresh_extras_logger):
monkeypatch.setenv("LITELLM_LOG", "ERROR")
reloaded = importlib.reload(extras_logging).logger
assert reloaded is fresh_extras_logger
assert reloaded.isEnabledFor(logging.INFO) is False
assert reloaded.isEnabledFor(logging.ERROR) is True
@pytest.mark.parametrize("litellm_log", [None, "info", "DEBUG"])
def test_unset_or_verbose_litellm_log_keeps_extras_info_lines(monkeypatch, fresh_extras_logger, litellm_log):
if litellm_log is None:
monkeypatch.delenv("LITELLM_LOG", raising=False)
else:
monkeypatch.setenv("LITELLM_LOG", litellm_log)
assert importlib.reload(extras_logging).logger.isEnabledFor(logging.INFO) is True

View file

@ -91,6 +91,154 @@ async def test_jwt_to_virtual_key_mapping_resolution():
prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called()
@pytest.mark.asyncio
async def test_colliding_claim_value_from_another_issuer_does_not_resolve_to_the_wrong_virtual_key():
"""LIT-7417: a mapping registered for one issuer must not answer a lookup from a
DIFFERENT issuer whose claim value happens to collide, even though both issuers
map the same claim field (``sub``) to a virtual key."""
issuer_a = "https://issuer-a.example.com"
issuer_b = "https://issuer-b.example.com"
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
virtual_key_claim_field="sub", virtual_key_mapping_cache_ttl=3600
)
rows = [
{
"jwt_issuer": issuer_b,
"jwt_claim_name": "sub",
"jwt_claim_value": "dev-alice",
"token": "hashed-issuer-b-key",
"is_active": True,
}
]
async def fake_find_first(where):
for row in rows:
if all(row.get(k) == v for k, v in where.items()):
return MagicMock(**row)
return None
prisma_client = MagicMock()
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(side_effect=fake_find_first)
# Dependency-inject the resolved key via the cache (IdentityStore._resolve_key
# reads it from here) instead of monkeypatching IdentityStore itself.
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key="hashed-issuer-b-key",
value=UserAPIKeyAuth(token="hashed-issuer-b-key", team_id="issuer-b-team"),
)
# The rightful owner: issuer-b's own claim resolves to its mapping.
owner_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_b, "sub": "dev-alice"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert isinstance(owner_result, UserAPIKeyAuth)
assert owner_result.token == "hashed-issuer-b-key"
# A validly-signed token from issuer-a carrying the SAME claim value must not
# inherit issuer-b's mapping. Default behavior is fallback_team_mapping, so a
# correctly-scoped miss returns None instead of resolving to issuer-b's key.
colliding_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_a, "sub": "dev-alice"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert colliding_result is None
@pytest.mark.asyncio
async def test_global_mapping_resolution_is_cached_under_the_global_key_not_the_requesting_issuer():
"""LIT-7417: caching a global (unscoped) mapping's hit under the REQUESTING
issuer's key would leave every issuer that falls back to it holding its own
stale copy after the row is updated/deleted -- CRUD only evicts the cache key
computed from the row's own scope (global), so a copy cached under some other
issuer's key would keep resolving to the old token until TTL. Caching it under
the global key instead means every issuer shares (and CRUD correctly evicts)
the exact same entry."""
issuer_a = "https://issuer-a.example.com"
issuer_b = "https://issuer-b.example.com"
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
issuers=[
{
"issuer": issuer_a,
"jwks_url": f"{issuer_a}/jwks",
"virtual_key_claim_field": "sub",
"disable_audience_validation": True,
},
{
"issuer": issuer_b,
"jwks_url": f"{issuer_b}/jwks",
"virtual_key_claim_field": "sub",
"disable_audience_validation": True,
},
]
)
rows = [
{
"jwt_issuer": "",
"jwt_claim_name": "sub",
"jwt_claim_value": "legacy-user",
"token": "hashed-legacy-key",
"is_active": True,
}
]
async def fake_find_first(where):
for row in rows:
if all(row.get(k) == v for k, v in where.items()):
return MagicMock(**row)
return None
prisma_client = MagicMock()
find_first = AsyncMock(side_effect=fake_find_first)
prisma_client.db.litellm_jwtkeymapping.find_first = find_first
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key="hashed-legacy-key",
value=UserAPIKeyAuth(token="hashed-legacy-key", team_id="legacy-team"),
)
resolved_a = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_a, "sub": "legacy-user"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert isinstance(resolved_a, UserAPIKeyAuth)
assert find_first.await_count == 2 # issuer-a-scoped miss, then global hit
# issuer-b resolving the SAME global mapping must hit the cache issuer-a's
# resolution populated, not issue a fresh DB query for the global row again.
resolved_b = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_b, "sub": "legacy-user"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert isinstance(resolved_b, UserAPIKeyAuth)
assert resolved_b.token == "hashed-legacy-key"
assert find_first.await_count == 3 # +1 for issuer-b's own issuer-scoped miss; global tier served from cache
@pytest.mark.asyncio
async def test_jwt_to_virtual_key_mapping_no_mapping():
"""
@ -223,6 +371,7 @@ def test_to_response_excludes_token():
now = datetime.now(timezone.utc)
mock_mapping = MagicMock()
mock_mapping.id = "mapping-1"
mock_mapping.jwt_issuer = None
mock_mapping.jwt_claim_name = "email"
mock_mapping.jwt_claim_value = "user@example.com"
mock_mapping.token = "hashed_secret_value"
@ -275,10 +424,12 @@ def _mock_mapping(
id="mapping-1",
claim_name="email",
claim_value="user@example.com",
issuer=None,
):
now = datetime.now(timezone.utc)
m = MagicMock()
m.id = id
m.jwt_issuer = issuer
m.jwt_claim_name = claim_name
m.jwt_claim_value = claim_value
m.token = "hashed_token"
@ -485,6 +636,35 @@ async def test_create_success_returns_response_without_token():
assert result.jwt_claim_name == "email"
@pytest.mark.asyncio
async def test_create_without_issuer_stores_empty_string_not_null():
"""LIT-7417: the DB column is NOT NULL (see schema.prisma). Storing a real NULL
for an unscoped mapping would let Postgres accept unlimited duplicate unscoped
rows for the same claim (NULL is never equal to NULL in a unique constraint),
so two mappings for the same claim value could point at two different keys with
no conflict, and resolution would pick whichever one Postgres returns first."""
from litellm.proxy._types import CreateJWTKeyMappingRequest
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.create.return_value = _mock_mapping()
mock_cache = AsyncMock()
data = CreateJWTKeyMappingRequest(jwt_claim_name="sub", jwt_claim_value="dev-alice", key="sk-test-key")
with (
patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point
"litellm.proxy.proxy_server.prisma_client", mock_prisma
),
patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point
"litellm.proxy.proxy_server.user_api_key_cache", mock_cache
),
):
await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
sent_data = mock_prisma.db.litellm_jwtkeymapping.create.call_args.kwargs["data"]
assert sent_data["jwt_issuer"] == ""
# ──────────────────────────────────────────────
# Tests: unregistered_jwt_client_behavior
# ──────────────────────────────────────────────

View file

@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
responses_reasoning_items_from_thinking_blocks,
split_concatenated_json_objects,
strip_encrypted_reasoning_from_messages,
system_messages_first,
update_messages_with_model_file_ids,
)
@ -1107,6 +1108,38 @@ def test_drop_tool_reference_parts_leaves_non_tool_messages_alone():
assert result[2]["content"] == ""
class TestSystemMessagesFirst:
def test_stable_partition_keeps_order_within_each_group(self):
messages = [
{"role": "user", "content": "u1"},
{"role": "system", "content": "s1"},
{"role": "assistant", "content": "a1"},
{"role": "developer", "content": "d1"},
{"role": "tool", "tool_call_id": "c1", "content": "t1"},
{"role": "system", "content": "s2"},
]
result = system_messages_first(messages)
assert [m["content"] for m in result] == ["s1", "d1", "s2", "u1", "a1", "t1"]
assert [m["content"] for m in messages] == ["u1", "s1", "a1", "d1", "t1", "s2"]
assert all(
result_message is original for result_message, original in zip(result[3:], messages[::2], strict=True)
)
@pytest.mark.parametrize(
"messages",
[
[],
[{"role": "user", "content": "u1"}, {"role": "assistant", "content": "a1"}],
[{"role": "system", "content": "s1"}, {"role": "user", "content": "u1"}],
[{"role": "system", "content": "s1"}, {"role": "system", "content": "s2"}],
],
)
def test_already_ordered_messages_come_back_unchanged(self, messages):
assert system_messages_first(messages) == messages
class TestFlattenTopLevelSchemaCombinators:
def _customer_anyof_schema(self):
return {

View file

@ -2,6 +2,7 @@ import base64
import json
import logging
import os
import re
from typing import Final
from unittest.mock import MagicMock, patch
@ -19,6 +20,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
_convert_to_bedrock_tool_call_invoke,
_convert_to_bedrock_tool_call_result,
anthropic_messages_pt,
convert_to_anthropic_tool_result,
convert_to_gemini_tool_call_result,
make_valid_bedrock_tool_name,
ollama_pt,
@ -2208,6 +2210,104 @@ def test_bedrock_tool_call_invoke_empty_arguments():
assert result[0]["toolUse"]["input"] == {}
_BEDROCK_TOOL_USE_ID_RE = re.compile(r"^[a-zA-Z0-9_.:-]{1,64}$")
@pytest.mark.parametrize(
"tool_call_id",
[
"call_" + "x" * 100,
"call|with|pipes",
"call_" + "y" * 60 + "|end",
"call:ok.dots-and_under",
"",
],
)
def test_bedrock_tool_use_id_is_sanitized_consistently_for_invoke_and_result(tool_call_id):
"""
Regression test for https://github.com/BerriAI/litellm/issues/34239: client-minted
tool_call ids longer than 64 chars or with chars outside [a-zA-Z0-9_.:-] made Bedrock
return a 400. The invoke and result paths must produce the same valid toolUseId so the
toolUse/toolResult pair still correlates.
"""
invoke = _convert_to_bedrock_tool_call_invoke(
[
{
"id": tool_call_id,
"type": "function",
"function": {"name": "get_weather", "arguments": '{"location": "Boston"}'},
}
]
)
result = _convert_to_bedrock_tool_call_result(
{"tool_call_id": tool_call_id, "role": "tool", "name": "get_weather", "content": "sunny"}
)
tool_use_id = invoke[0]["toolUse"]["toolUseId"]
assert _BEDROCK_TOOL_USE_ID_RE.match(tool_use_id)
assert result["toolResult"]["toolUseId"] == tool_use_id
def test_bedrock_tool_use_id_valid_ids_pass_through_unchanged():
result = _convert_to_bedrock_tool_call_result(
{"tool_call_id": "tooluse_Ab.c:1-2_3", "role": "tool", "name": "f", "content": "ok"}
)
assert result["toolResult"]["toolUseId"] == "tooluse_Ab.c:1-2_3"
def test_bedrock_tool_use_id_truncation_keeps_distinct_ids_distinct():
prefix = "call_" + "z" * 70
ids = {
_convert_to_bedrock_tool_call_result(
{"tool_call_id": f"{prefix}{suffix}", "role": "tool", "name": "f", "content": "ok"}
)["toolResult"]["toolUseId"]
for suffix in ("a", "b")
}
assert len(ids) == 2
assert all(len(i) == 64 for i in ids)
def test_bedrock_tool_use_id_replaced_chars_do_not_collide_with_existing_ids():
ids = {
_convert_to_bedrock_tool_call_result({"tool_call_id": i, "role": "tool", "name": "f", "content": "ok"})[
"toolResult"
]["toolUseId"]
for i in ("call|x", "call_x")
}
assert len(ids) == 2
def test_bedrock_tool_call_invoke_concatenated_json_long_id_stays_within_limit():
long_id = "call_" + "q" * 62
result = _convert_to_bedrock_tool_call_invoke(
[
{
"id": long_id,
"type": "function",
"function": {"name": "run", "arguments": '{"cmd":"a"}{"cmd":"b"}'},
}
]
)
ids = [block["toolUse"]["toolUseId"] for block in result]
assert len(ids) == 2
assert len(set(ids)) == 2
assert all(_BEDROCK_TOOL_USE_ID_RE.match(i) for i in ids)
@pytest.mark.parametrize(
("tool_call_id", "expected"),
[
("call|with|pipes", "call_with_pipes"),
("call:ok.dots", "call_ok_dots"),
("call_" + "x" * 100, "call_" + "x" * 100),
("toolu_01AbC-xyz", "toolu_01AbC-xyz"),
("", "tool_use_id"),
],
)
def test_anthropic_tool_use_id_keeps_pattern_only_rewrite_with_no_cap_or_hash(tool_call_id, expected):
result = convert_to_anthropic_tool_result({"role": "tool", "tool_call_id": tool_call_id, "content": "ok"})
assert result["tool_use_id"] == expected
def test_bedrock_tool_call_invoke_concatenated_json():
"""
Tool call whose arguments contain multiple concatenated JSON objects

View file

@ -135,9 +135,7 @@ def test_fill_missing_requires_per_rule_opt_in(restore_generalizations):
"supports_vision": True,
}
restore_generalizations(
[{"name": "base", "pattern": r"^acme-", "model_info": {"supports_reasoning": True}}]
)
restore_generalizations([{"name": "base", "pattern": r"^acme-", "model_info": {"supports_reasoning": True}}])
assert match_fill_missing_generalizations("acme-1", "openai") is None
restore_generalizations(
@ -451,6 +449,94 @@ def shipped_cost_map(monkeypatch):
set_fallback_generalizations(previous_rules)
@pytest.mark.parametrize(
"model,provider",
[
("gemini-4-pro", "gemini"),
("gemini/gemini-4-pro", None),
("gemini-3.9-flash-lite-preview-09-2026", "vertex_ai"),
("vertex_ai/gemini-4-pro", None),
("gemini-4-pro-preview-customtools", "gemini"),
("google/gemini-4-pro", "openrouter"),
("google/gemini-4-pro", "deepinfra"),
("google/gemini-4-pro", "vercel_ai_gateway"),
("google.gemini-4-pro", "oci"),
("databricks-gemini-4-1-pro", "databricks"),
],
)
def test_shipped_gemini_chat_baseline_resolves_unmapped_ids(shipped_cost_map, model, provider):
assert model not in litellm.model_cost
if provider == "gemini":
assert f"gemini/{model}" not in litellm.model_cost
elif provider in {"openrouter", "deepinfra", "vercel_ai_gateway", "oci", "databricks"}:
assert f"{provider}/{model}" not in litellm.model_cost
info = litellm.get_model_info(model, custom_llm_provider=provider)
assert info["litellm_provider"] == (provider or model.split("/")[0])
assert info["mode"] == "chat"
assert not info.get("max_input_tokens")
assert info["supports_reasoning"] is True
assert info["supports_function_calling"] is True
assert info["supports_tool_choice"] is True
assert info["supports_system_messages"] is True
assert info["supports_vision"] is True
assert info["supports_response_schema"] is True
assert info["supports_pdf_input"] is True
assert info["supports_prompt_caching"] is True
assert info["supports_web_search"] is True
assert not info.get("input_cost_per_token")
assert not info.get("output_cost_per_token")
def test_shipped_gemini_chat_baseline_loses_to_perplexity_exact_entries(shipped_cost_map):
info = litellm.get_model_info("google/gemini-2.5-pro", custom_llm_provider="perplexity")
entry = litellm.model_cost["perplexity/google/gemini-2.5-pro"]
assert info["mode"] == "responses"
assert entry["supports_reasoning"] is False
def test_shipped_gemini_chat_baseline_skips_non_chat_and_pre_2_5_ids(shipped_cost_map):
for model in (
"gemini/gemini-4-flash-image",
"gemini/gemini-3.9-flash-preview-tts",
"gemini/gemini-4-flash-live-preview",
"gemini/gemini-4-flash-native-audio",
"gemini/gemini-embedding-4",
"gemini/gemini-2.5-computer-use-preview-12-2026",
"gemini/gemini-2.0-flash-new",
"gemini/gemini-1.5-pro-new",
"gemini/gemini-4-flashy",
"gemini/gemini-4-flash-transcribe",
"gemini/gemini-4-flash-live-translate-preview",
"databricks-gemini-3-1-flash-image",
"openrouter/google/gemini-2.0-flash-001",
):
assert match_capability_generalizations(model) is None, model
def test_shipped_gemini_chat_baseline_keeps_reasoning_effort_on_unmapped_model(shipped_cost_map):
assert litellm.supports_reasoning(model="gemini-4-pro", custom_llm_provider="gemini") is True
optional_params = litellm.utils.get_optional_params(
model="gemini-4-pro",
custom_llm_provider="gemini",
reasoning_effort="medium",
drop_params=False,
)
assert isinstance(optional_params, dict)
assert optional_params["thinkingConfig"]["thinkingBudget"] > 0
assert optional_params["thinkingConfig"]["includeThoughts"] is True
def test_shipped_gemini_chat_baseline_loses_to_exact_entries(shipped_cost_map):
model = "gemini-2.5-flash-lite"
info = litellm.get_model_info(model, custom_llm_provider="gemini")
entry = litellm.model_cost["gemini/gemini-2.5-flash-lite"]
assert info["max_tokens"] == entry["max_tokens"]
assert info["input_cost_per_token"] == entry["input_cost_per_token"]
assert entry["input_cost_per_token"] > 0
def test_shipped_bare_claude_id_routes_to_anthropic(shipped_cost_map):
_, provider, _, _ = litellm.get_llm_provider(model="claude-haiku-4-6")
assert provider == "anthropic"

View file

@ -1,7 +1,5 @@
import pytest
from litellm.litellm_core_utils.get_supported_openai_params import (
get_supported_openai_params,
)
@ -33,9 +31,7 @@ def test_base_model_label_alone_lacks_bedrock_tools():
"""The label by itself does not advertise tools; this is what made the union
necessary. Guards against the discrepancy disappearing (and the regression test
above silently passing for the wrong reason)."""
params = get_supported_openai_params(
model=BEDROCK_LABEL, custom_llm_provider="bedrock"
)
params = get_supported_openai_params(model=BEDROCK_LABEL, custom_llm_provider="bedrock")
assert params is not None
assert "tools" not in params
@ -46,14 +42,8 @@ def test_base_model_is_additive_not_replacement():
Bedrock: real id supports ``tools`` but not the label's reasoning hint; the union
must contain the real model's ``tools`` regardless of the label being a subset."""
real_only = set(
get_supported_openai_params(
model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock"
)
)
label_only = set(
get_supported_openai_params(model=BEDROCK_LABEL, custom_llm_provider="bedrock")
)
real_only = set(get_supported_openai_params(model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock"))
label_only = set(get_supported_openai_params(model=BEDROCK_LABEL, custom_llm_provider="bedrock"))
combined = set(
get_supported_openai_params(
model=BEDROCK_REAL_MODEL,
@ -70,19 +60,15 @@ def test_base_model_is_additive_not_replacement():
def test_base_model_adds_capabilities_the_real_model_lacks():
"""Regression for #27717 (the behavior the union must preserve).
``gemini-3.1-pro`` isn't in the cost map so it advertises no reasoning support,
``gemini-exp-9999`` isn't in the cost map so it advertises no reasoning support,
but the registered ``gemini-3.1-pro-preview`` base_model does. The hint must add
``reasoning_effort``/``thinking`` without the call erroring."""
real_only = set(
get_supported_openai_params(
model="gemini-3.1-pro", custom_llm_provider="gemini"
)
)
real_only = set(get_supported_openai_params(model="gemini-exp-9999", custom_llm_provider="gemini"))
assert "reasoning_effort" not in real_only
combined = set(
get_supported_openai_params(
model="gemini-3.1-pro",
model="gemini-exp-9999",
custom_llm_provider="gemini",
base_model="gemini-3.1-pro-preview",
)
@ -93,21 +79,15 @@ def test_base_model_adds_capabilities_the_real_model_lacks():
def test_no_base_model_is_unchanged():
"""Omitting ``base_model`` must resolve purely from ``model``."""
with_none = get_supported_openai_params(
model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock", base_model=None
)
plain = get_supported_openai_params(
model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock"
)
with_none = get_supported_openai_params(model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock", base_model=None)
plain = get_supported_openai_params(model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock")
assert with_none == plain
def test_base_model_equal_to_model_is_unchanged():
"""A ``base_model`` identical to ``model`` must not double-resolve or reorder."""
plain = get_supported_openai_params(
model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock"
)
plain = get_supported_openai_params(model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock")
same = get_supported_openai_params(
model=BEDROCK_REAL_MODEL,
custom_llm_provider="bedrock",
@ -152,14 +132,10 @@ def test_bedrock_converse_alias_resolves_like_bedrock():
params saw no Bedrock capabilities for a Converse model invoked via the alias."""
anthropic_model = "bedrock/converse/us.anthropic.claude-sonnet-4-6"
via_alias = get_supported_openai_params(
model=anthropic_model, custom_llm_provider="bedrock_converse"
)
via_alias = get_supported_openai_params(model=anthropic_model, custom_llm_provider="bedrock_converse")
assert via_alias is not None
assert via_alias == get_supported_openai_params(
model=anthropic_model, custom_llm_provider="bedrock"
)
assert via_alias == get_supported_openai_params(model=anthropic_model, custom_llm_provider="bedrock")
assert "web_search_options" not in via_alias
assert "tools" in via_alias
@ -167,9 +143,7 @@ def test_bedrock_converse_alias_resolves_like_bedrock():
def test_bedrock_converse_alias_keeps_nova_web_search_options():
"""Nova on the ``bedrock_converse`` alias still advertises web_search_options, proving the
alias routes through the model-aware config rather than a blanket Bedrock default."""
nova_params = get_supported_openai_params(
model="amazon.nova-pro-v1:0", custom_llm_provider="bedrock_converse"
)
nova_params = get_supported_openai_params(model="amazon.nova-pro-v1:0", custom_llm_provider="bedrock_converse")
assert nova_params is not None
assert "web_search_options" in nova_params

View file

@ -6554,9 +6554,9 @@ async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_o
assert pre_choice["metadata"]["litellm_gateway_injected_cache"] == ""
def _responses_ws_logging_obj() -> LitellmLogging:
def _responses_ws_logging_obj(model: str = "gpt-4o") -> LitellmLogging:
return LitellmLogging(
model="gpt-4o",
model=model,
messages=[],
stream=False,
call_type=CallTypes.aresponses_websocket.value,
@ -6638,6 +6638,62 @@ def test_normalize_logging_result_bills_incomplete_responses_websocket_turns():
assert normalized.usage.total_tokens == 75
def test_normalize_logging_result_prices_responses_websocket_at_returned_service_tier():
"""Issue #41299: a WebSocket turn billed at priority tier reported it on
response.completed.response.service_tier, but the logging object dropped it and the
session was priced at the default tier."""
events = [
{"type": "response.created", "response": {}},
{
"type": "response.completed",
"response": {
"service_tier": "priority",
"usage": {"input_tokens": 100, "output_tokens": 40, "total_tokens": 140},
},
},
]
normalized = _responses_ws_logging_obj(model="gpt-5.4").normalize_logging_result(result=events)
assert isinstance(normalized, LiteLLMRealtimeStreamLoggingObject)
assert normalized.service_tier == "priority"
usage = ResponseAPIUsage(input_tokens=100, output_tokens=40, total_tokens=140)
ws_cost = litellm.completion_cost(
completion_response=normalized,
model="gpt-5.4",
call_type=CallTypes.aresponses_websocket.value,
custom_llm_provider="openai",
)
priority_http_cost = litellm.completion_cost(
completion_response=ResponsesAPIResponse(
id="resp-priority",
created_at=1700000000,
output=[],
service_tier="priority",
usage=usage,
),
model="gpt-5.4",
call_type=CallTypes.aresponses.value,
custom_llm_provider="openai",
)
default_http_cost = litellm.completion_cost(
completion_response=ResponsesAPIResponse(
id="resp-default",
created_at=1700000000,
output=[],
service_tier="default",
usage=usage,
),
model="gpt-5.4",
call_type=CallTypes.aresponses.value,
custom_llm_provider="openai",
)
assert ws_cost == priority_http_cost
assert priority_http_cost > default_http_cost
def test_get_standard_logging_object_payload_reads_overhead_from_logging_obj_for_dict_results(logging_obj):
"""LIT-5466: /v1/messages returns a plain dict with no _hidden_params, so the overhead
recorded on the logging object must reach hidden_params.litellm_overhead_time_ms (SpendLogs)."""

View file

@ -132,6 +132,32 @@ def test_transform_request_drops_tool_reference_parts():
assert request["messages"][2]["content"] == ""
@pytest.mark.parametrize(
"enabled, expected", [(False, ("hi", "sys", "reply", "more")), (True, ("sys", "hi", "reply", "more"))]
)
def test_transform_request_system_messages_first_follows_global_flag(monkeypatch, enabled, expected):
"""Azure OpenAI shares OpenAI's prefix-matched prompt cache, so the same flag moves
system messages ahead of the conversation on the Azure request body."""
monkeypatch.setattr(litellm, "openai_system_messages_first", enabled)
messages = [
{"role": "user", "content": "hi"},
{"role": "system", "content": "sys"},
{"role": "assistant", "content": "reply"},
{"role": "user", "content": "more"},
]
request = AzureOpenAIConfig().transform_request(
model="gpt-4o",
messages=messages,
optional_params={},
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
assert tuple(m["content"] for m in request["messages"]) == expected
assert [m["content"] for m in messages] == ["hi", "sys", "reply", "more"]
@pytest.mark.parametrize(
"model, emitted_key, absent_key",
[

View file

@ -68,3 +68,24 @@ def test_azure_o_series_transform_request_flattens_top_level_anyof():
assert parameters["required"] == ["id"]
assert "anyOf" in tool["function"]["parameters"]
assert optional_params["tools"][0] is tool
def test_azure_o_series_transform_request_moves_system_messages_first(monkeypatch):
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
messages = [
{"role": "user", "content": "hi"},
{"role": "developer", "content": "dev"},
{"role": "assistant", "content": "reply"},
{"role": "user", "content": "more"},
]
request = AzureOpenAIO1Config().transform_request(
model="o3-mini",
messages=messages,
optional_params={},
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
assert [m["content"] for m in request["messages"]] == ["dev", "hi", "reply", "more"]
assert [m["content"] for m in messages] == ["hi", "dev", "reply", "more"]

View file

@ -430,6 +430,25 @@ class TestGeminiVideoConfig:
assert result.usage["video_resolution"] == "1080p"
assert result.usage["duration_seconds"] == 8.0
def test_transform_video_create_response_usage_includes_video_count(self):
"""Regression for LIT-6896: sampleCount (number of generated videos) is copied into usage for billing."""
mock_response = Mock(spec=httpx.Response)
mock_response.json.return_value = {"name": "operations/generate_1234567890"}
request_data = {
"instances": [{"prompt": "Test"}],
"parameters": {"durationSeconds": 8, "sampleCount": 3},
}
result = self.config.transform_video_create_response(
model="gemini/veo-3.1-fast-generate-preview",
raw_response=mock_response,
logging_obj=self.mock_logging_obj,
custom_llm_provider="gemini",
request_data=request_data,
)
assert result.usage is not None
assert result.usage["video_count"] == 3
assert result.usage["duration_seconds"] == 8.0
def test_transform_video_create_response_cost_tracking_with_different_durations(
self,
):

View file

@ -0,0 +1,296 @@
import json
from types import MappingProxyType
import httpx
import pytest
import litellm
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.llms.nvidia_nim.passthrough.transformation import (
NvidiaNimPassthroughConfig,
nvidia_nim_model_group_in_path,
nvidia_nim_model_groups,
nvidia_nim_router_model_in_endpoint,
)
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
NIM_BASE = "http://nim.internal:8000"
INFER_BODY = {
"input": [
{"type": "image_url", "url": "data:image/png;base64,AAAA"},
{"type": "image_url", "url": "data:image/png;base64,BBBB"},
]
}
@pytest.fixture(autouse=True)
def clear_nvidia_nim_env(monkeypatch):
for env_var in ("NVIDIA_NIM_API_BASE", "NVIDIA_NIM_API_KEY"):
monkeypatch.delenv(env_var, raising=False)
monkeypatch.setattr(litellm, "api_base", None)
monkeypatch.setattr(litellm, "api_key", None)
def test_provider_config_manager_resolves_nvidia_nim_passthrough_config():
config = ProviderConfigManager.get_provider_passthrough_config(
model="nvidia/nemoretriever-page-elements-v2", provider=LlmProviders.NVIDIA_NIM
)
assert isinstance(config, NvidiaNimPassthroughConfig)
@pytest.mark.parametrize(
"api_base, endpoint, litellm_params, expected",
[
(NIM_BASE, "nim-page/v1/infer", {"litellm_metadata": {"model_group": "nim-page"}}, f"{NIM_BASE}/v1/infer"),
(
f"{NIM_BASE}/v1",
"nim-page/v1/infer",
{"litellm_metadata": {"model_group": "nim-page"}},
f"{NIM_BASE}/v1/infer",
),
(f"{NIM_BASE}/v1/", "/v1/infer", {}, f"{NIM_BASE}/v1/infer"),
(NIM_BASE, "v1/infer", {}, f"{NIM_BASE}/v1/infer"),
(f"{NIM_BASE}/v2", "v1/infer", {}, f"{NIM_BASE}/v2/v1/infer"),
(f"{NIM_BASE}/infer", "infer", {}, f"{NIM_BASE}/infer/infer"),
(NIM_BASE, "nvidia/nemoretriever-page-elements-v2/v1/infer", {}, f"{NIM_BASE}/v1/infer"),
(
NIM_BASE,
"nvidia/nemoretriever-page-elements-v2/v1/infer",
{"litellm_metadata": {"model_group": "nvidia"}},
f"{NIM_BASE}/v1/infer",
),
],
)
def test_relay_url_strips_the_model_group_and_never_doubles_the_api_version(
api_base, endpoint, litellm_params, expected
):
url, base = NvidiaNimPassthroughConfig().get_complete_url(
api_base=api_base,
api_key=None,
model="nvidia/nemoretriever-page-elements-v2",
endpoint=endpoint,
request_query_params=None,
litellm_params=litellm_params,
)
assert str(url) == expected
assert base == expected.removesuffix("/v1/infer").removesuffix("/infer")
def test_query_params_are_forwarded_on_the_relay_url():
url, _ = NvidiaNimPassthroughConfig().get_complete_url(
api_base=NIM_BASE,
api_key=None,
model="nvidia/nemoretriever-page-elements-v2",
endpoint="v1/infer",
request_query_params={"timeout": "30"},
litellm_params={},
)
assert str(url) == f"{NIM_BASE}/v1/infer?timeout=30"
def test_env_api_base_is_used_when_the_deployment_has_none(monkeypatch):
monkeypatch.setenv("NVIDIA_NIM_API_BASE", f"{NIM_BASE}/v1")
url, _ = NvidiaNimPassthroughConfig().get_complete_url(
api_base=None,
api_key=None,
model="nvidia/nemoretriever-page-elements-v2",
endpoint="v1/infer",
request_query_params=None,
litellm_params={},
)
assert str(url) == f"{NIM_BASE}/v1/infer"
def test_missing_api_base_raises_instead_of_building_a_relative_url():
with pytest.raises(ValueError, match="NVIDIA_NIM_API_BASE"):
NvidiaNimPassthroughConfig().get_complete_url(
api_base=None,
api_key=None,
model="nvidia/nemoretriever-page-elements-v2",
endpoint="v1/infer",
request_query_params=None,
litellm_params={},
)
def test_deployment_key_becomes_a_bearer_token_and_caller_headers_are_kept():
caller_headers = MappingProxyType({"x-request-id": "abc"})
headers = NvidiaNimPassthroughConfig().validate_environment(
headers=caller_headers,
model="nvidia/nemoretriever-page-elements-v2",
messages=[],
optional_params={},
litellm_params={},
api_key="nvapi-secret",
)
assert headers == {"x-request-id": "abc", "Authorization": "Bearer nvapi-secret"}
def test_self_hosted_nim_without_a_key_sends_no_authorization_header():
headers = NvidiaNimPassthroughConfig().validate_environment(
headers={}, model="nvidia/x", messages=[], optional_params={}, litellm_params={}, api_key=None
)
assert "Authorization" not in headers
def test_env_api_key_fills_in_when_the_deployment_has_none(monkeypatch):
monkeypatch.setenv("NVIDIA_NIM_API_KEY", "nvapi-from-env")
assert NvidiaNimPassthroughConfig.get_api_key(None) == "nvapi-from-env"
assert NvidiaNimPassthroughConfig.get_api_key("nvapi-deployment") == "nvapi-deployment"
@pytest.mark.parametrize(
"endpoint, router_models, expected",
[
("nim-page/v1/infer", ("nim-page", "nim-table"), "nim-page"),
("/nim-page/v1/infer", ("nim-page",), "nim-page"),
(
"nvidia/nemoretriever-page-elements-v2/v1/infer",
("nvidia/nemoretriever-page-elements-v2",),
"nvidia/nemoretriever-page-elements-v2",
),
("nim/v1/infer", ("nim", "nim/v1"), "nim/v1"),
("v1/infer", ("nim-page",), None),
("nim-page-elements/v1/infer", ("nim-page",), None),
("", ("nim-page",), None),
],
)
def test_router_model_in_endpoint_takes_the_longest_leading_model_group(endpoint, router_models, expected):
assert nvidia_nim_router_model_in_endpoint(endpoint, frozenset(router_models)) == expected
def _deployment(model_name: str, model: str, custom_llm_provider: str | None = None):
litellm_params = (
{"model": model}
if custom_llm_provider is None
else {"model": model, "custom_llm_provider": custom_llm_provider}
)
return {"model_name": model_name, "litellm_params": litellm_params}
MIXED_DEPLOYMENTS = (
_deployment("nim-page", "nvidia_nim/nvidia/nemoretriever-page-elements-v2"),
_deployment("nim-table", "nvidia/nemoretriever-table-structure-v1", custom_llm_provider="nvidia_nim"),
_deployment("mixed", "nvidia_nim/nvidia/nemoretriever-page-elements-v2"),
_deployment("mixed", "openai/gpt-4o"),
_deployment("gpt-4o", "openai/gpt-4o"),
)
def test_model_groups_only_admit_groups_whose_every_deployment_is_nim_backed():
assert nvidia_nim_model_groups(MIXED_DEPLOYMENTS) == frozenset({"nim-page", "nim-table"})
assert nvidia_nim_model_groups(None) == frozenset()
@pytest.mark.parametrize(
"path, expected",
[
("/nvidia_nim/nim-page/v1/infer", "nim-page"),
("/NVIDIA_NIM/nim-table/v1/infer", "nim-table"),
("nim-page/v1/infer", "nim-page"),
("/nvidia_nim/mixed/v1/infer", None),
("mixed/v1/infer", None),
("/nvidia_nim/gpt-4o/v1/infer", None),
("/nvidia_nim/v1/infer", None),
],
)
def test_model_group_in_path_resolves_the_same_nim_only_groups_for_routes_and_endpoints(path, expected):
assert nvidia_nim_model_group_in_path(path, MIXED_DEPLOYMENTS) == expected
@pytest.mark.parametrize("request_data, expected", [({"stream": True}, True), ({"stream": False}, False), ({}, False)])
def test_is_streaming_request_reads_the_stream_flag(request_data, expected):
assert NvidiaNimPassthroughConfig().is_streaming_request("v1/infer", request_data) is expected
def test_non_streaming_relay_logs_the_upstream_json_body():
response = httpx.Response(
200,
json={"data": [{"index": 0, "bounding_boxes": {}}]},
request=httpx.Request("POST", f"{NIM_BASE}/v1/infer"),
)
result = NvidiaNimPassthroughConfig().logging_non_streaming_response(
model="nvidia/nemoretriever-page-elements-v2",
custom_llm_provider="nvidia_nim",
httpx_response=response,
request_data=INFER_BODY,
logging_obj=None, # pyright: ignore[reportArgumentType] # not read for a plain passthrough body
endpoint="v1/infer",
)
assert result == {"response": {"data": [{"index": 0, "bounding_boxes": {}}]}}
@pytest.mark.asyncio
async def test_object_detection_relay_sends_the_native_body_unchanged_to_v1_infer():
upstream_requests: list[httpx.Request] = []
def nim(request: httpx.Request) -> httpx.Response:
upstream_requests.append(request)
return httpx.Response(200, json={"data": [{"index": 0}, {"index": 1}]}, headers={"x-nim": "1"})
client = AsyncHTTPHandler()
client.client = httpx.AsyncClient(transport=httpx.MockTransport(nim))
response = await litellm.allm_passthrough_route(
model="nvidia_nim/nvidia/nemoretriever-page-elements-v2",
endpoint="nim-page/v1/infer",
method="POST",
api_base=f"{NIM_BASE}/v1",
api_key="nvapi-secret",
json=dict(INFER_BODY),
litellm_metadata={"model_group": "nim-page"},
client=client,
)
(sent,) = upstream_requests
assert str(sent.url) == f"{NIM_BASE}/v1/infer"
assert json.loads(sent.content) == INFER_BODY
assert sent.headers["authorization"] == "Bearer nvapi-secret"
assert response.status_code == 200
assert response.headers["x-nim"] == "1"
assert response.json() == {"data": [{"index": 0}, {"index": 1}]}
@pytest.mark.asyncio
async def test_router_relay_reaches_v1_infer_when_the_group_name_is_a_leading_segment_of_the_model_id():
upstream_requests: list[httpx.Request] = []
def nim(request: httpx.Request) -> httpx.Response:
upstream_requests.append(request)
return httpx.Response(200, json={"data": [{"index": 0}]})
client = AsyncHTTPHandler()
client.client = httpx.AsyncClient(transport=httpx.MockTransport(nim))
router = litellm.Router(
model_list=[
{
"model_name": "nvidia",
"litellm_params": {
"model": "nvidia_nim/nvidia/nemoretriever-page-elements-v2",
"api_base": NIM_BASE,
"api_key": "nvapi-secret",
},
}
]
)
response = await router.allm_passthrough_route(
model="nvidia", endpoint="nvidia/v1/infer", method="POST", json=dict(INFER_BODY), client=client
)
(sent,) = upstream_requests
assert str(sent.url) == f"{NIM_BASE}/v1/infer"
assert json.loads(sent.content) == INFER_BODY
assert response.status_code == 200

View file

@ -1124,6 +1124,80 @@ class TestToolReferenceStripping:
assert request["messages"][2]["content"] == ""
class TestSystemMessagesFirst:
"""With litellm.openai_system_messages_first on, requests bound for OpenAI put system and
developer messages ahead of the conversation, keeping each group's order, so the instruction
prefix stays byte-stable for OpenAI's prefix-matched prompt cache."""
MESSAGES: Final = (
{"role": "user", "content": "first turn"},
{"role": "system", "content": "sys 1"},
{"role": "assistant", "content": "reply"},
{"role": "developer", "content": "dev"},
{"role": "user", "content": "second turn"},
{"role": "system", "content": "sys 2"},
)
ORIGINAL_ORDER: Final = ("first turn", "sys 1", "reply", "dev", "second turn", "sys 2")
ORDERED: Final = ("sys 1", "dev", "sys 2", "first turn", "reply", "second turn")
def setup_method(self):
self.config = OpenAIGPTConfig()
def _messages(self):
return [dict(m) for m in self.MESSAGES]
def _transform(self, provider):
return self.config.transform_request(
model="gpt-4.1",
messages=self._messages(),
optional_params={},
litellm_params={"custom_llm_provider": provider},
headers={},
)
def test_default_off_keeps_caller_order(self, monkeypatch):
monkeypatch.setattr(litellm, "openai_system_messages_first", False)
assert tuple(m["content"] for m in self._transform("openai")["messages"]) == self.ORIGINAL_ORDER
def test_moves_system_and_developer_messages_first_for_openai(self, monkeypatch):
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
assert tuple(m["content"] for m in self._transform("openai")["messages"]) == self.ORDERED
def test_leaves_openai_compatible_providers_alone(self, monkeypatch):
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
assert tuple(m["content"] for m in self._transform("deepseek")["messages"]) == self.ORIGINAL_ORDER
def test_does_not_mutate_caller_messages(self, monkeypatch):
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
messages = self._messages()
self.config.transform_request(
model="gpt-4.1",
messages=messages,
optional_params={},
litellm_params={"custom_llm_provider": "openai"},
headers={},
)
assert tuple(m["content"] for m in messages) == self.ORIGINAL_ORDER
@pytest.mark.asyncio
async def test_async_transform_request_moves_system_messages_first(self, monkeypatch):
class UninstantiatedOpenAIGPTConfig(OpenAIGPTConfig):
_is_base_class = True
def __init__(self) -> None:
pass
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
request = await UninstantiatedOpenAIGPTConfig().async_transform_request(
model="gpt-4.1",
messages=self._messages(),
optional_params={},
litellm_params={"custom_llm_provider": "openai"},
headers={},
)
assert tuple(m["content"] for m in request["messages"]) == self.ORDERED
class TestOpenAIPromptCacheBreakpointChatPath:
"""Chat-path shape for OpenAI explicit prompt caching (#37509)."""

View file

@ -9,7 +9,95 @@ import litellm
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import (
VertexPassthroughLoggingHandler,
)
from litellm.types.utils import PassthroughCallTypes
from litellm.types.utils import ModelResponse, PassthroughCallTypes
_OMNI_INTERACTIONS_USAGE: Final = {
"total_tokens": 4041,
"total_input_tokens": 12,
"input_tokens_by_modality": [{"modality": "text", "tokens": 12}],
"total_output_tokens": 4009,
"output_tokens_by_modality": [
{"modality": "text", "tokens": 9},
{"modality": "video", "tokens": 4000},
],
"total_tool_use_tokens": 0,
"total_thought_tokens": 20,
}
def test_interactions_create_response_logs_modality_usage_and_cost() -> None:
"""
Regression for LIT-6896: gemini-omni Interactions passthrough rows were logged
with zero tokens and zero spend. Input, text-output and video-output tokens
must land in usage, priced with the model's per-modality rates, and the
response id must stay the litellm_call_id so SpendLogs keep their request_id.
"""
logging_obj = MagicMock()
logging_obj.model_call_details = {}
logging_obj.optional_params = {}
logging_obj.litellm_call_id = "call-6896"
response = httpx.Response(
status_code=200,
json={
"id": "interactions/abc",
"model": "gemini-omni-flash-preview",
"status": "completed",
"outputs": [{"type": "text", "text": "hi"}],
"usage": _OMNI_INTERACTIONS_USAGE,
},
)
result = VertexPassthroughLoggingHandler.vertex_passthrough_handler(
httpx_response=response,
logging_obj=logging_obj,
url_route="https://aiplatform.googleapis.com/v1beta1/projects/p/locations/global/interactions",
result=response.text,
start_time=datetime.now(),
end_time=datetime.now(),
cache_hit=False,
request_body={"model": "gemini-omni-flash-preview", "input": [{"type": "text", "text": "say hi"}]},
)
model_response = result["result"]
assert isinstance(model_response, ModelResponse)
assert model_response.id == "call-6896"
usage = model_response.usage
assert usage.prompt_tokens == 12
assert usage.completion_tokens == 4009 + 20
assert usage.completion_tokens_details.text_tokens == 9
assert usage.completion_tokens_details.video_tokens == 4000
model_info = litellm.get_model_info(model="gemini-omni-flash-preview", custom_llm_provider="vertex_ai")
expected_cost = (
12 * model_info["input_cost_per_token"]
+ (9 + 20) * model_info["output_cost_per_token"]
+ 4000 * model_info["output_cost_per_video_token"]
)
assert result["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert result["kwargs"]["custom_llm_provider"] == "vertex_ai"
assert logging_obj.model_call_details["model"] == "gemini-omni-flash-preview"
assert logging_obj.model_call_details["custom_llm_provider"] == "vertex_ai"
def test_interactions_response_without_usage_falls_back_to_generic_logging() -> None:
logging_obj = MagicMock()
logging_obj.model_call_details = {}
logging_obj.optional_params = {}
response = httpx.Response(status_code=200, json={"id": "interactions/abc", "status": "in_progress"})
result = VertexPassthroughLoggingHandler.vertex_passthrough_handler(
httpx_response=response,
logging_obj=logging_obj,
url_route="https://aiplatform.googleapis.com/v1beta1/projects/p/locations/global/interactions",
result=response.text,
start_time=datetime.now(),
end_time=datetime.now(),
cache_hit=False,
request_body={"agent": "projects/p/locations/global/reasoningEngines/1"},
)
assert result["result"] is None
assert "response_cost" not in result["kwargs"]
def test_lyria_predict_response_preserves_audio_response_and_logs_cost(

View file

@ -717,6 +717,33 @@ class TestVertexAIVideoConfig:
assert video_obj.usage["duration_seconds"] == 8.0
assert video_obj.usage["video_resolution"] == "1080p"
@pytest.mark.parametrize(
"sample_count,expected_video_count",
[(2, 2), (1, 1), (None, None), (0, None), ("2", None)],
ids=["two", "one", "unset", "zero", "string"],
)
def test_transform_video_create_response_usage_includes_video_count(self, sample_count, expected_video_count):
"""Regression for LIT-6896: sampleCount is the number of generated videos and must reach usage for billing."""
mock_response = Mock(spec=httpx.Response)
mock_response.json.return_value = {
"name": "projects/p/locations/us-central1/publishers/google/models/veo-3.1-fast-generate-001/operations/op-1"
}
parameters = {"durationSeconds": 4, "resolution": "720p"}
if sample_count is not None:
parameters["sampleCount"] = sample_count
video_obj = self.config.transform_video_create_response(
model="veo-3.1-fast-generate-001",
raw_response=mock_response,
logging_obj=self.mock_logging_obj,
custom_llm_provider="vertex_ai",
request_data={"instances": [{"prompt": "a red ball"}], "parameters": parameters},
)
assert video_obj.usage is not None
assert video_obj.usage["duration_seconds"] == 4.0
assert video_obj.usage.get("video_count") == expected_video_count
def test_transform_video_remix_request_not_supported(self):
"""Test that video remix raises NotImplementedError."""
with pytest.raises(NotImplementedError, match="Video remix is not supported"):

View file

@ -22,6 +22,7 @@ from litellm.proxy.auth.auth_utils import (
get_key_mcp_rpm_limit,
get_key_model_rpm_limit,
get_key_model_tpm_limit,
get_key_own_model_rate_limit,
get_key_tag_rpm_limit,
get_model_from_request,
get_project_model_rpm_limit,
@ -141,6 +142,35 @@ class TestLogOnceIfBudgetReservationDisabled:
class TestGetKeyModelRpmLimit:
"""Tests for get_key_model_rpm_limit function."""
def test_own_limit_excludes_team_metadata(self):
"""A team-only limit is inherited, not owned: the key resolves it but does not override it."""
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-123",
metadata={"some_other_key": "value"},
team_metadata={"model_rpm_limit": {"gpt-4": 50}, "model_tpm_limit": {"gpt-4": 500}},
)
assert get_key_model_rpm_limit(user_api_key_dict) == {"gpt-4": 50}
assert get_key_own_model_rate_limit(user_api_key_dict, "model_rpm_limit") is None
assert get_key_own_model_rate_limit(user_api_key_dict, "model_tpm_limit") is None
def test_own_limit_resolves_metadata_then_model_max_budget(self):
from_metadata = UserAPIKeyAuth(
api_key="sk-123",
metadata={"model_rpm_limit": {"gpt-4": 100}},
model_max_budget={"gpt-4": {"rpm_limit": 10, "tpm_limit": 1000}},
team_metadata={"model_rpm_limit": {"gpt-4": 50}},
)
assert get_key_own_model_rate_limit(from_metadata, "model_rpm_limit") == {"gpt-4": 100}
assert get_key_own_model_rate_limit(from_metadata, "model_tpm_limit") == {"gpt-4": 1000}
from_budget = UserAPIKeyAuth(
api_key="sk-123",
model_max_budget={"gpt-4": {"rpm_limit": 10}, "gpt-3.5-turbo": {"tpm_limit": 1000}},
team_metadata={"model_rpm_limit": {"gpt-4": 50}},
)
assert get_key_own_model_rate_limit(from_budget, "model_rpm_limit") == {"gpt-4": 10}
assert get_key_own_model_rate_limit(from_budget, "model_tpm_limit") == {"gpt-3.5-turbo": 1000}
def test_returns_key_metadata_when_present(self):
"""Key metadata takes priority over team metadata."""
user_api_key_dict = UserAPIKeyAuth(
@ -823,6 +853,82 @@ def test_get_model_from_request_azure_relay_routes_use_the_model_group_in_the_pa
assert get_model_from_request(request_data=request_data, route=route, llm_router=_azure_relay_router()) == expected
def _nvidia_nim_relay_router():
from litellm.router import Router
return Router(
model_list=[
{
"model_name": "nim-page-elements",
"litellm_params": {
"model": "nvidia_nim/nvidia/nemoretriever-page-elements-v2",
"api_base": "http://nim-a.internal:8000",
"api_key": "k",
},
},
{
"model_name": "nvidia/nemoretriever-table-structure-v1",
"litellm_params": {
"model": "nvidia_nim/nvidia/nemoretriever-table-structure-v1",
"api_base": "http://nim-b.internal:8000",
"api_key": "k",
},
},
{
"model_name": "gpt-4o",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "k"},
},
{
"model_name": "detect",
"litellm_params": {
"model": "nvidia_nim/nvidia/nemoretriever-page-elements-v2",
"api_base": "http://nim-a.internal:8000",
"api_key": "k",
},
},
{
"model_name": "detect",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "k"},
},
]
)
NIM_INFER_BODY = {"input": [{"type": "image_url", "url": "data:image/png;base64,AAAA"}]}
@pytest.mark.parametrize(
"route, request_data, expected",
[
("/nvidia_nim/nim-page-elements/v1/infer", NIM_INFER_BODY, "nim-page-elements"),
(
"/nvidia_nim/nim-page-elements/v1/infer",
{"model": "nvidia/nemoretriever-table-structure-v1"},
"nim-page-elements",
),
(
"/nvidia_nim/nvidia/nemoretriever-table-structure-v1/v1/infer",
NIM_INFER_BODY,
"nvidia/nemoretriever-table-structure-v1",
),
("/nvidia_nim/v1/infer", NIM_INFER_BODY, None),
("/nvidia_nim/unknown-group/v1/infer", NIM_INFER_BODY, None),
("/nvidia_nim/nim-page-elements-v2/v1/infer", NIM_INFER_BODY, None),
("/nvidia_nim/gpt-4o/v1/infer", NIM_INFER_BODY, None),
("/nvidia_nim/detect/v1/infer", NIM_INFER_BODY, None),
],
)
def test_get_model_from_request_nvidia_nim_relay_routes_use_the_model_group_in_the_path(route, request_data, expected):
assert (
get_model_from_request(request_data=request_data, route=route, llm_router=_nvidia_nim_relay_router())
== expected
)
def test_get_model_from_request_nvidia_nim_relay_without_a_router_has_no_model():
assert get_model_from_request(request_data=NIM_INFER_BODY, route="/nvidia_nim/nim-page-elements/v1/infer") is None
def test_get_model_from_request_includes_file_endpoint_header_model():
assert (
get_model_from_request(

View file

@ -693,6 +693,7 @@ def test_virtual_key_allowed_routes_with_litellm_routes_member_name_denied():
"/anthropic/v1/count_tokens",
"/gemini/v1/models",
"/gemini/countTokens",
"/nvidia_nim/nim-page-elements/v1/infer",
],
)
def test_virtual_key_llm_api_route_includes_passthrough_prefix(route):

View file

@ -6,6 +6,7 @@ import subprocess
import sys
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from functools import partial
from pathlib import Path
from textwrap import dedent
from types import SimpleNamespace
@ -32,8 +33,15 @@ from litellm.proxy._types import (
JWTRoutingOverride,
)
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.auth_checks import TeamNotFoundError, UserNotFoundError, get_key_object, _cache_key_object
from litellm.proxy.auth.auth_checks import (
TeamNotFoundError,
UserNotFoundError,
get_key_object,
_cache_key_object,
jwt_key_mapping_cache_key,
)
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model
from litellm.proxy.auth.user_api_key_auth import (
_check_key_model_budget_with_fallback,
_ensure_litellm_received_at_on_request_state,
@ -7948,13 +7956,38 @@ def _per_issuer_virtual_key_jwt_handler(
def _fake_prisma_with_jwt_key_mapping(hashed_token: str | None) -> tuple[SimpleNamespace, AsyncMock]:
"""Every ``find_first`` call (issuer-scoped or global fallback) resolves the same way."""
find_first = AsyncMock(return_value=None if hashed_token is None else SimpleNamespace(token=hashed_token))
prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first)))
return prisma_client, find_first
def _mapping_where(claim_name: str, claim_value: str) -> dict[str, str | bool]:
return {"jwt_claim_name": claim_name, "jwt_claim_value": claim_value, "is_active": True}
def _fake_prisma_jwt_key_mapping_table(rows: list[dict[str, object]]) -> tuple[SimpleNamespace, AsyncMock]:
"""A ``find_first`` whose result depends on the ``where`` clause, like a real table.
Matches a row when every key present in ``where`` equals that key on the row --
a key ``get_jwt_key_mapping_object`` omits (e.g. old, issuer-blind code never
sending ``jwt_issuer``) does not constrain the match, exactly like Prisma.
"""
async def _find_first(where: dict[str, object]) -> SimpleNamespace | None:
for row in rows:
if all(row.get(k) == v for k, v in where.items()):
return SimpleNamespace(**row)
return None
find_first = AsyncMock(side_effect=_find_first)
prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first)))
return prisma_client, find_first
def _mapping_where(claim_name: str, claim_value: str, jwt_issuer: str | None) -> dict[str, str | bool]:
return {
"jwt_claim_name": claim_name,
"jwt_claim_value": claim_value,
"jwt_issuer": jwt_issuer or "",
"is_active": True,
}
@pytest.mark.asyncio
@ -7978,11 +8011,13 @@ async def test_per_issuer_virtual_key_claim_field_selects_the_issuer_mapping_for
proxy_logging_obj=MagicMock(),
)
find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7"))
# Issuer-scoped lookup hits on the first query, so no global fallback query runs.
find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7", ISSUER_TWO))
assert isinstance(resolved, UserAPIKeyAuth)
assert resolved.token == "hashed-mapped-key"
assert resolved.team_id == "svc-team"
assert await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:svc-account-7") == "hashed-mapped-key"
cache_key = jwt_key_mapping_cache_key("sub", "svc-account-7", ISSUER_TWO)
assert await user_api_key_cache.async_get_cache(cache_key) == "hashed-mapped-key"
@pytest.mark.asyncio
@ -8015,7 +8050,11 @@ async def test_per_issuer_reject_behavior_does_not_leak_into_the_team_issuer():
assert exc.value.status_code == 403
assert "No registered mapping for sub='unknown-svc'" in str(exc.value.detail)
find_first.assert_awaited_once_with(where=_mapping_where("sub", "unknown-svc"))
# REJECT checks the issuer-scoped row first, then falls back to a global (NULL-issuer) row.
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
_mapping_where("sub", "unknown-svc", ISSUER_TWO),
_mapping_where("sub", "unknown-svc", None),
]
@pytest.mark.asyncio
@ -8025,7 +8064,10 @@ async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_rej
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub", global_behavior="auto_register")
prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None)
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(key="jwt_key_mapping:sub:admin-7", value=_JWT_PROXY_ADMIN_SENTINEL)
# Sentinel cached under issuer-one's own key -- must never answer issuer-two's lookup.
await user_api_key_cache.async_set_cache(
key=jwt_key_mapping_cache_key("sub", "admin-7", ISSUER_ONE), value=_JWT_PROXY_ADMIN_SENTINEL
)
auto_register_issuer_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "admin-7"},
@ -8050,7 +8092,10 @@ async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_rej
assert exc.value.status_code == 403
assert "No registered mapping for sub='admin-7'" in str(exc.value.detail)
find_first.assert_awaited_once_with(where=_mapping_where("sub", "admin-7"))
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
_mapping_where("sub", "admin-7", ISSUER_TWO),
_mapping_where("sub", "admin-7", None),
]
@pytest.mark.asyncio
@ -8079,7 +8124,125 @@ async def test_issuer_without_virtual_key_claim_field_falls_back_to_the_global_f
assert with_claim is None
assert without_claim is None
find_first.assert_awaited_once_with(where=_mapping_where("client_id", "app-9"))
# without_claim has no claim value and returns before ever reaching the DB.
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
_mapping_where("client_id", "app-9", ISSUER_ONE),
_mapping_where("client_id", "app-9", None),
]
@pytest.mark.asyncio
async def test_colliding_claim_value_from_another_issuer_does_not_resolve_to_the_wrong_virtual_key():
"""LIT-7417: a mapping registered for one issuer must not answer a lookup from a
DIFFERENT issuer whose claim value happens to collide, even though both issuers
use the same claim field (``sub``) for their virtual-key mapping."""
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
prisma_client, find_first = _fake_prisma_jwt_key_mapping_table(
[
{
"jwt_issuer": ISSUER_TWO,
"jwt_claim_name": "sub",
"jwt_claim_value": "dev-alice",
"token": "hashed-issuer-b-key",
"is_active": True,
}
]
)
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key="hashed-issuer-b-key",
value=UserAPIKeyAuth(token="hashed-issuer-b-key", api_key="hashed-issuer-b-key", team_id="issuer-b-team"),
)
owner_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "dev-alice"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert isinstance(owner_result, UserAPIKeyAuth)
assert owner_result.token == "hashed-issuer-b-key"
# issuer-one's behavior is fallback_team_mapping: a correctly-scoped miss must
# return None (fall through to team-based JWT auth), never issuer-two's key.
colliding_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "dev-alice"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert colliding_result is None
assert find_first.await_count == 3 # owner hit (1 call) + colliding miss (issuer-scoped + global fallback)
@pytest.mark.asyncio
async def test_cached_resolution_for_one_issuer_does_not_leak_to_a_colliding_issuer():
"""A cached positive resolution must be keyed by issuer too, or a colliding
claim value from another issuer could be served straight from cache without
ever reaching the (correctly issuer-scoped) DB lookup."""
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
prisma_client, find_first = _fake_prisma_jwt_key_mapping_table([])
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key=jwt_key_mapping_cache_key("sub", "dev-alice", ISSUER_TWO), value="hashed-issuer-b-key"
)
colliding_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "dev-alice"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert colliding_result is None
# Must have gone to the DB rather than serving issuer-two's cached token.
assert find_first.await_count == 2
@pytest.mark.asyncio
async def test_issuer_agnostic_mapping_matches_every_issuer():
"""A mapping created before issuer scoping existed (``jwt_issuer`` is NULL) keeps
matching any issuer, so existing global mappings are not broken by this fix."""
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
prisma_client, _find_first = _fake_prisma_jwt_key_mapping_table(
[
{
"jwt_issuer": "",
"jwt_claim_name": "sub",
"jwt_claim_value": "legacy-user",
"token": "hashed-legacy-key",
"is_active": True,
}
]
)
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key="hashed-legacy-key",
value=UserAPIKeyAuth(token="hashed-legacy-key", api_key="hashed-legacy-key", team_id="legacy-team"),
)
for issuer in (ISSUER_ONE, ISSUER_TWO):
resolved = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer, "sub": "legacy-user"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert isinstance(resolved, UserAPIKeyAuth)
assert resolved.token == "hashed-legacy-key"
@pytest.mark.asyncio
@ -8137,3 +8300,189 @@ async def test_auth_flow_enters_virtual_key_mapping_when_only_an_issuer_configur
assert resolve_mock.await_args.kwargs["jwt_claims"][JWTHandler.LITELLM_JWT_ISSUER_CLAIM] == ISSUER_TWO
assert result.api_key == "hashed-mapped-key"
assert result.team_id == "svc-team"
def _alias_router() -> litellm.Router:
return litellm.Router(
model_list=[
{"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}}
for name in ("claude-haiku", "claude-sonnet")
]
)
def _alias_request(route: str, data: dict, content_type: str = "application/json", path_params: dict | None = None):
"""A request as auth sees it: the body already read once and cached alongside its parsed form."""
from starlette.requests import Request
scope = {
"type": "http",
"method": "POST",
"path": route,
"headers": [(b"content-type", content_type.encode())],
"query_string": b"",
"path_params": path_params or {},
"parsed_body": (tuple(data), data),
}
request = Request(scope)
request._body = json.dumps(data).encode()
return request
async def _enforce_alias_access(token: UserAPIKeyAuth, data: dict, route: str, request, router: litellm.Router):
from litellm.proxy.auth.user_api_key_auth import _enforce_key_and_fallback_model_access
await _enforce_key_and_fallback_model_access(
valid_token=token,
request_data=data,
route=route,
request=request,
llm_model_list=router.model_list,
llm_router=router,
)
def _alias_token(monkeypatch, level: str, alias: dict, models: list) -> UserAPIKeyAuth:
"""A key whose ``router_settings.model_group_alias`` lives on the key itself or on its cached team row."""
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
if level == "key":
return UserAPIKeyAuth(models=models, router_settings={"model_group_alias": alias})
cache = UserApiKeyCache()
team = LiteLLM_TeamTableCachedObj(team_id="team-alias", models=models, router_settings={"model_group_alias": alias})
cache.set_cache(key="team_id:team-alias", value=team)
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", MagicMock())
return UserAPIKeyAuth(team_id="team-alias", models=models)
@pytest.mark.asyncio
@pytest.mark.parametrize("level", ["key", "team"])
@pytest.mark.parametrize(
"route",
["/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/openai/v1/responses", "/cursor/chat/completions"],
)
async def test_router_settings_model_group_alias_authorizes_target_for_key(monkeypatch, level, route):
"""LIT-3054: a key allowed only the alias target must be able to call the alias, and a key not
allowed the target must still be denied even when the alias itself is what it requested."""
router = _alias_router()
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router)
data = {"model": "AgentX-LLM", "messages": [{"role": "user", "content": "hi"}]}
request = _alias_request(route, data)
token = _alias_token(monkeypatch, level, {"AgentX-LLM": "claude-haiku"}, ["claude-haiku"])
await _enforce_alias_access(token, data, route, request, router)
assert data["model"] == "claude-haiku"
assert (await request.json())["model"] == "claude-haiku"
assert json.loads(await request.body())["model"] == "claude-haiku"
assert request.scope["parsed_body"][1]["model"] == "claude-haiku"
assert get_client_requested_model(request) == "AgentX-LLM"
denied = _alias_token(monkeypatch, level, {"AgentX-LLM": "claude-sonnet"}, ["claude-haiku"])
denied_data = {"model": "AgentX-LLM"}
with pytest.raises(ProxyException) as exc:
await _enforce_alias_access(denied, denied_data, route, _alias_request(route, denied_data), router)
assert "claude-sonnet" in exc.value.message
@pytest.mark.asyncio
async def test_router_settings_model_group_alias_leaves_form_bodies_alone(monkeypatch):
"""LIT-3054: a multipart body cannot be re-serialized as JSON, so auth must not rewrite it."""
router = _alias_router()
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router)
data = {"model": "AgentX-LLM"}
route = "/v1/audio/transcriptions"
request = _alias_request(route, data, content_type="multipart/form-data; boundary=x")
token = _alias_token(monkeypatch, "key", {"AgentX-LLM": "claude-haiku"}, ["claude-haiku", "AgentX-LLM"])
await _enforce_alias_access(token, data, route, request, router)
assert data["model"] == "AgentX-LLM"
assert get_client_requested_model(request) is None
@pytest.mark.asyncio
async def test_router_settings_model_group_alias_rewrite_keeps_query_params_out_of_body(monkeypatch):
"""LIT-3054: auth merges query params into its own copy of the body; the rewrite must not forward them."""
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body, populate_request_with_path_params
router = _alias_router()
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router)
body = {"model": "AgentX-LLM", "messages": [{"role": "user", "content": "hi"}]}
request = _alias_request("/v1/chat/completions", body)
request.scope["query_string"] = b"api-version=2024-10-21&stream=true"
data = populate_request_with_path_params(request_data=await _read_request_body(request), request=request)
assert data["api-version"] == "2024-10-21"
token = _alias_token(monkeypatch, "key", {"AgentX-LLM": "claude-haiku"}, ["claude-haiku"])
await _enforce_alias_access(token, data, "/v1/chat/completions", request, router)
downstream = await _read_request_body(request)
assert downstream == {**body, "model": "claude-haiku"}
assert json.loads(await request.body()) == downstream
assert await request.json() == downstream
def _user_defined_pass_through_endpoint():
from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER
async def endpoint():
return None
setattr(endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True)
return endpoint
@pytest.mark.asyncio
@pytest.mark.parametrize("user_defined", [False, True])
async def test_router_settings_model_group_alias_leaves_pass_through_bodies_alone(monkeypatch, user_defined):
"""LIT-3054: pass-through handlers forward the body verbatim to the provider, so auth must not rewrite it.
Built-in provider handlers bind ``{endpoint:path}``; user-defined ones carry the pass-through marker."""
router = _alias_router()
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router)
data = {"model": "AgentX-LLM", "messages": [{"role": "user", "content": "hi"}]}
route = "/custom-upstream/chat" if user_defined else "/anthropic/v1/messages"
request = _alias_request(route, data, path_params={} if user_defined else {"endpoint": "v1/messages"})
if user_defined:
request.scope["endpoint"] = _user_defined_pass_through_endpoint()
LiteLLMRoutes.openai_routes.value.append(route)
token = _alias_token(monkeypatch, "key", {"AgentX-LLM": "claude-haiku"}, ["claude-haiku", "AgentX-LLM"])
try:
await _enforce_alias_access(token, data, route, request, router)
finally:
if user_defined:
LiteLLMRoutes.openai_routes.value.remove(route)
assert data["model"] == "AgentX-LLM"
assert (await request.json())["model"] == "AgentX-LLM"
assert get_client_requested_model(request) is None
@pytest.mark.asyncio
@pytest.mark.parametrize("target, expect_denied", [("claude-haiku", False), ("claude-sonnet", True)])
async def test_router_settings_model_group_alias_authorizes_target_for_team(monkeypatch, target, expect_denied):
"""LIT-3054: the team allowlist check in common_checks must judge the alias target, not the alias."""
import litellm.proxy.proxy_server as _proxy_server_mod
from litellm.proxy.auth.user_api_key_auth import _authorize_authenticated_request
router = _alias_router()
logging_obj = MagicMock(post_call_failure_hook=AsyncMock(return_value=None))
attrs = {**_proxy_attrs_for_centralized_checks(), "llm_router": router, "proxy_logging_obj": logging_obj}
for k, v in attrs.items():
monkeypatch.setattr(_proxy_server_mod, k, v)
token = _alias_token(monkeypatch, "team", {"AgentX-LLM": target}, ["claude-haiku"])
token.team_models = ["claude-haiku"]
data = {"model": "AgentX-LLM", "messages": [{"role": "user", "content": "hi"}]}
route = "/v1/chat/completions"
request = _alias_request(route, data)
authorize = partial(
_authorize_authenticated_request,
user_api_key_auth_obj=token,
request=request,
request_data=data,
route=route,
api_key="sk-test",
)
if expect_denied:
with pytest.raises(ProxyException) as exc:
await authorize()
assert exc.value.type == ProxyErrorTypes.team_model_access_denied
assert target in exc.value.message
return
await authorize()
assert (await request.json())["model"] == target
assert get_client_requested_model(request) == "AgentX-LLM"

View file

@ -1,7 +1,9 @@
import inspect
import json
import os
import subprocess
import sys
from pathlib import Path
from unittest.mock import patch
import click
@ -9,10 +11,9 @@ import pytest
import requests
from click.testing import CliRunner
from litellm.proxy.client.cli.commands.agents import (
AgentRunError,
ModelSyncArgs,
ModelSyncSkipped,
_hand_off,
_replace_process,
@ -22,6 +23,7 @@ from litellm.proxy.client.cli.commands.agents import (
agent_model_sync_env,
agent_profile,
build_agent_env,
codex_model_sync_args,
opencode_model_sync_env,
run_agent,
verify_proxy_key,
@ -55,6 +57,112 @@ class _Recorder:
return self.returns
_STOCK_REASONING_LEVELS = [
{"effort": "low", "description": "Fast responses with lighter reasoning"},
{"effort": "medium", "description": "Balances speed and reasoning depth for everyday tasks"},
{"effort": "high", "description": "Greater reasoning depth for complex problems"},
]
_STOCK_MODELS = {
"gpt-5.6-terra": {
"slug": "gpt-5.6-terra",
"display_name": "GPT-5.6 Terra",
"description": "Balanced agentic coding model for everyday work.",
"default_reasoning_level": "medium",
"supported_reasoning_levels": _STOCK_REASONING_LEVELS,
"shell_type": "unified_exec",
"visibility": "list",
"supported_in_api": True,
"priority": 7,
"availability_nux": None,
"upgrade": None,
"base_instructions": "You are Codex, a coding agent based on GPT-5.6.",
"apply_patch_tool_type": "freeform",
"supports_parallel_tool_calls": True,
"context_window": 272000,
"comp_hash": "terra-hash",
},
"gpt-5.5": {
"slug": "gpt-5.5",
"display_name": "GPT-5.5",
"description": "Frontier model for complex coding, research, and real-world work.",
"default_reasoning_level": "medium",
"supported_reasoning_levels": _STOCK_REASONING_LEVELS,
"shell_type": "unified_exec",
"visibility": "list",
"supported_in_api": True,
"priority": 12,
"availability_nux": None,
"upgrade": None,
"base_instructions": "You are Codex, a coding agent based on GPT-5.",
"apply_patch_tool_type": "freeform",
"supports_parallel_tool_calls": True,
"context_window": 272000,
"comp_hash": "gpt-5.5-hash",
},
"gpt-5.4": {
"slug": "gpt-5.4",
"display_name": "GPT-5.4",
"description": "Strong model for everyday coding.",
"default_reasoning_level": "medium",
"supported_reasoning_levels": _STOCK_REASONING_LEVELS,
"shell_type": "unified_exec",
"visibility": "hide",
"supported_in_api": True,
"priority": 16,
"availability_nux": None,
"upgrade": {
"model": "gpt-5.6-terra",
"migration_markdown": "GPT-5.4 is no longer available. Switch to GPT-5.6 Terra to continue.",
"retirement_at": "2026-08-31T19:00:00Z",
},
"base_instructions": "You are Codex, a coding agent based on GPT-5.",
"apply_patch_tool_type": "freeform",
"supports_parallel_tool_calls": True,
"context_window": 272000,
"comp_hash": "gpt-5.4-hash",
},
"codex-auto-review": {
"slug": "codex-auto-review",
"display_name": "Codex Auto Review",
"description": None,
"supported_reasoning_levels": [],
"shell_type": "unified_exec",
"visibility": "hide",
"supported_in_api": False,
"priority": 43,
"availability_nux": None,
"upgrade": None,
"base_instructions": "You are Codex, reviewing a change.",
"apply_patch_tool_type": None,
"supports_parallel_tool_calls": True,
"context_window": 272000,
"comp_hash": "review-hash",
},
}
_STOCK_CATALOG = json.dumps({"models": list(_STOCK_MODELS.values())})
class _FakeRun:
"""A `codex` that prints `stock` from a bare `debug models` and answers a catalog override with `returncode`.
`stock=None` is a Codex with no `debug models` at all: every call answers with `returncode` and `stderr`.
"""
def __init__(self, returncode=0, stderr="", stock=_STOCK_CATALOG):
self.returncode = returncode
self.stderr = stderr
self.stock = stock
self.calls = []
def __call__(self, args, **kwargs):
self.calls.append((args, kwargs))
if self.stock is not None and "model_catalog_json=" not in str(args):
return subprocess.CompletedProcess(args, 0, self.stock, "")
return subprocess.CompletedProcess(args, self.returncode, "", self.stderr)
class _FakeJsonResponse:
def __init__(self, status_code, payload=None):
self.status_code = status_code
@ -90,9 +198,7 @@ class TestAgentProfile:
class TestBuildAgentEnv:
def test_anthropic_profile_uses_bare_root_and_bearer(self):
env = build_agent_env(
{}, "http://localhost:4000/", "sk-key", frozenset({"anthropic"})
)
env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"anthropic"}))
assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000"
assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key"
assert env["ENABLE_TOOL_SEARCH"] == "true"
@ -128,9 +234,7 @@ class TestBuildAgentEnv:
assert "ANTHROPIC_API_KEY" not in env
def test_openai_profile_appends_v1(self):
env = build_agent_env(
{}, "http://localhost:4000/", "sk-key", frozenset({"openai"})
)
env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"openai"}))
assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1"
assert env["OPENAI_API_KEY"] == "sk-key"
assert "ANTHROPIC_BASE_URL" not in env
@ -138,9 +242,7 @@ class TestBuildAgentEnv:
assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in env
def test_both_profiles_set_everything(self):
env = build_agent_env(
{}, "http://localhost:4000", "sk-key", frozenset({"anthropic", "openai"})
)
env = build_agent_env({}, "http://localhost:4000", "sk-key", frozenset({"anthropic", "openai"}))
assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000"
assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1"
assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key"
@ -148,9 +250,7 @@ class TestBuildAgentEnv:
assert env["ENABLE_TOOL_SEARCH"] == "true"
def test_litellm_profile_exports_only_the_proxy_key(self):
env = build_agent_env(
{}, "http://localhost:4000/", "sk-key", frozenset({"litellm"})
)
env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"litellm"}))
assert env["LITELLM_PROXY_API_KEY"] == "sk-key"
assert "ANTHROPIC_BASE_URL" not in env
assert "OPENAI_BASE_URL" not in env
@ -158,9 +258,7 @@ class TestBuildAgentEnv:
def test_preserves_unrelated_env_and_does_not_mutate_input(self):
base = {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"}
env = build_agent_env(
base, "http://localhost:4000", "sk-key", frozenset({"anthropic"})
)
env = build_agent_env(base, "http://localhost:4000", "sk-key", frozenset({"anthropic"}))
assert env["PATH"] == "/usr/bin"
assert base == {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"}
@ -322,9 +420,7 @@ class TestOpencodeModelSync:
assert "refused" in result.reason
def test_non_200_is_reported(self):
result = opencode_model_sync_env(
{}, "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500)
)
result = opencode_model_sync_env({}, "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500))
assert isinstance(result, ModelSyncSkipped)
assert "HTTP 500" in result.reason
@ -335,10 +431,10 @@ class TestOpencodeModelSync:
assert isinstance(result, ModelSyncSkipped)
assert "unexpected body" in result.reason
@pytest.mark.parametrize("command", ["claude", "codex", "/usr/bin/claude"])
def test_only_opencode_syncs(self, command):
@pytest.mark.parametrize("command", ["claude", "pi", "/usr/bin/claude"])
def test_only_opencode_and_codex_sync(self, command):
def boom(*a, **k):
raise AssertionError("no agent other than opencode should call the proxy")
raise AssertionError("no agent other than opencode or codex should call the proxy")
assert agent_model_sync_env(command, {}, "http://localhost:4000", "sk-key", False, get=boom) == {}
@ -367,7 +463,369 @@ class TestOpencodeModelSync:
assert _default_of(opencode_model_sync_env, "get") is requests.get
class TestCodexModelSync:
@staticmethod
def _listing(*models):
return {"object": "list", "data": list(models)}
@staticmethod
def _row(model_id, **extra):
return {"id": model_id, "object": "model", "created": 1, "owned_by": "openai", **extra}
def _sync(self, listing, codex_home, base_url="http://localhost:4000/", run=None):
captured = {}
def fake_get(url, headers, timeout):
captured["url"] = url
captured["headers"] = headers
return _FakeResponse(200, listing)
result = codex_model_sync_args(
{"CODEX_HOME": str(codex_home)},
base_url,
"sk-key",
get=fake_get,
run=_FakeRun() if run is None else run,
)
return captured, result
@staticmethod
def _catalog_path(result):
assert isinstance(result, ModelSyncArgs)
flag, override = result.args
assert flag == "-c"
key, _, value = override.partition("=")
assert key == "model_catalog_json"
return json.loads(value)
def test_writes_catalog_under_codex_home_and_points_codex_at_it(self, tmp_path):
listing = self._listing(self._row("gpt-5.5", mode="chat"), self._row("claude-opus-4-7"))
captured, result = self._sync(listing, tmp_path / "codex")
assert captured["url"] == "http://localhost:4000/v1/models"
assert captured["headers"] == {"Authorization": "Bearer sk-key"}
path = self._catalog_path(result)
assert path == str(tmp_path / "codex" / "litellm-models.json")
text = (tmp_path / "codex" / "litellm-models.json").read_text()
assert "sk-key" not in text
catalog = json.loads(text)
assert [m["slug"] for m in catalog["models"]] == ["gpt-5.5", "claude-opus-4-7"]
assert [m["display_name"] for m in catalog["models"]] == ["GPT-5.5", "claude-opus-4-7"]
assert [m["priority"] for m in catalog["models"]] == [0, 1]
def _entries(self, codex_home):
return {m["slug"]: m for m in json.loads((codex_home / "litellm-models.json").read_text())["models"]}
def test_known_model_keeps_the_installed_codex_entry(self, tmp_path):
self._sync(self._listing(self._row("gpt-5.5", mode="chat")), tmp_path)
assert self._entries(tmp_path)["gpt-5.5"] == {**_STOCK_MODELS["gpt-5.5"], "priority": 0}
def test_hidden_stock_model_is_listed_when_the_proxy_serves_it(self, tmp_path):
self._sync(self._listing(self._row("gpt-5.4")), tmp_path)
entry = self._entries(tmp_path)["gpt-5.4"]
assert entry["visibility"] == "list"
assert entry["upgrade"] is None
assert entry["supported_reasoning_levels"] == _STOCK_REASONING_LEVELS
def test_api_disabled_stock_model_is_selectable_when_the_proxy_serves_it(self, tmp_path):
self._sync(self._listing(self._row("codex-auto-review")), tmp_path)
entry = self._entries(tmp_path)["codex-auto-review"]
assert entry["supported_in_api"] is True
assert entry["visibility"] == "list"
assert entry["base_instructions"] == _STOCK_MODELS["codex-auto-review"]["base_instructions"]
def test_stock_upgrade_nudge_survives_when_its_target_is_listed(self, tmp_path):
self._sync(self._listing(self._row("gpt-5.4"), self._row("gpt-5.6-terra")), tmp_path)
entries = self._entries(tmp_path)
assert entries["gpt-5.4"]["upgrade"] == _STOCK_MODELS["gpt-5.4"]["upgrade"]
assert [entries["gpt-5.4"]["priority"], entries["gpt-5.6-terra"]["priority"]] == [0, 1]
def test_stock_catalog_is_decoded_as_utf8_regardless_of_locale(self, tmp_path):
description = "Modelo equilibrado para el trabajo diario, con acentos y ñ."
catalog = {"models": [{**_STOCK_MODELS["gpt-5.5"], "description": description}]}
stock = json.dumps(catalog, ensure_ascii=False).encode("utf-8")
def locale_bound_run(args, **kwargs):
if "model_catalog_json=" in str(args):
return subprocess.CompletedProcess(args, 0, "", "")
return subprocess.CompletedProcess(args, 0, stock.decode(kwargs.get("encoding") or "ascii"), "")
_, result = self._sync(self._listing(self._row("gpt-5.5")), tmp_path, run=locale_bound_run)
assert isinstance(result, ModelSyncArgs)
written = json.loads((tmp_path / "litellm-models.json").read_text(encoding="utf-8"))["models"]
assert [m["description"] for m in written] == [description]
def test_unparseable_stock_catalog_is_reported(self, tmp_path):
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=_FakeRun(stock="not json"))
assert isinstance(result, ModelSyncSkipped)
assert result.reason.startswith("`codex debug models` printed no model catalog: ")
assert not (tmp_path / "litellm-models.json").exists()
def test_unknown_model_gets_the_fields_codex_requires(self, tmp_path):
_, result = self._sync(self._listing(self._row("m")), tmp_path)
entry = json.loads((tmp_path / "litellm-models.json").read_text())["models"][0]
assert entry["visibility"] == "list"
assert entry["supported_in_api"] is True
assert entry["shell_type"] == "unified_exec"
assert entry["supported_reasoning_levels"] == []
assert entry["truncation_policy"] == {"mode": "bytes", "limit": 10000}
assert entry["experimental_supported_tools"] == []
assert entry["support_verbosity"] is False
assert entry["supports_reasoning_summaries"] is False
assert entry["supports_parallel_tool_calls"] is False
for nullable in ("description", "availability_nux", "upgrade", "default_verbosity", "apply_patch_tool_type"):
assert nullable in entry and entry[nullable] is None
assert entry["base_instructions"].startswith("You are a coding agent running in the Codex CLI")
def test_context_window_comes_from_max_input_tokens_for_unknown_models_only(self, tmp_path):
listing = self._listing(
self._row("big", max_input_tokens=400000),
self._row("unknown"),
self._row("gpt-5.5", max_input_tokens=400000),
)
self._sync(listing, tmp_path)
models = self._entries(tmp_path)
assert models["big"]["context_window"] == 400000
assert models["unknown"]["context_window"] is None
assert models["gpt-5.5"]["context_window"] == 272000
def test_non_chat_models_are_left_out(self, tmp_path):
listing = self._listing(
self._row("chat", mode="chat"),
self._row("resp", mode="responses"),
self._row("embed", mode="embedding"),
self._row("img", mode="image_generation"),
)
self._sync(listing, tmp_path)
slugs = {m["slug"] for m in json.loads((tmp_path / "litellm-models.json").read_text())["models"]}
assert slugs == {"chat", "resp"}
def test_listing_without_chat_models_is_skipped_and_writes_nothing(self, tmp_path):
_, result = self._sync(self._listing(self._row("embed", mode="embedding")), tmp_path)
assert isinstance(result, ModelSyncSkipped)
assert "no chat models" in result.reason
assert not (tmp_path / "litellm-models.json").exists()
def test_catalog_is_rewritten_on_every_launch(self, tmp_path):
self._sync(self._listing(self._row("old")), tmp_path)
self._sync(self._listing(self._row("new")), tmp_path)
slugs = [m["slug"] for m in json.loads((tmp_path / "litellm-models.json").read_text())["models"]]
assert slugs == ["new"]
def test_catalog_is_replaced_whole_and_leaves_no_temp_files(self, tmp_path):
self._sync(self._listing(*(self._row(f"m{i}") for i in range(50))), tmp_path)
self._sync(self._listing(self._row("new")), tmp_path)
assert [p.name for p in tmp_path.iterdir()] == ["litellm-models.json"]
assert json.loads((tmp_path / "litellm-models.json").read_text())["models"][0]["slug"] == "new"
def test_defaults_to_dot_codex_in_home(self, tmp_path):
result = codex_model_sync_args(
{},
"http://localhost:4000",
"sk-key",
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
run=_FakeRun(),
home=lambda: tmp_path,
)
assert self._catalog_path(result) == str(tmp_path / ".codex" / "litellm-models.json")
def test_default_home_is_the_users(self):
assert _default_of(codex_model_sync_args, "home") == Path.home
def test_missing_base_instructions_is_reported_not_raised(self, tmp_path):
result = codex_model_sync_args(
{"CODEX_HOME": str(tmp_path)},
"http://localhost:4000",
"sk-key",
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
instructions_path=tmp_path / "missing.md",
)
assert isinstance(result, ModelSyncSkipped)
assert "could not read" in result.reason
assert not (tmp_path / "litellm-models.json").exists()
def test_unwritable_catalog_path_is_reported_not_raised(self, tmp_path):
blocker = tmp_path / "file"
blocker.write_text("")
_, result = self._sync(self._listing(self._row("m")), blocker / "codex")
assert isinstance(result, ModelSyncSkipped)
assert "could not write" in result.reason
def test_failed_replace_is_reported_and_leaves_no_temp_file(self, tmp_path):
(tmp_path / "litellm-models.json").mkdir()
_, result = self._sync(self._listing(self._row("m")), tmp_path)
assert isinstance(result, ModelSyncSkipped)
assert "could not write" in result.reason
assert [p.name for p in tmp_path.iterdir()] == ["litellm-models.json"]
def test_unreachable_proxy_is_reported_not_raised(self, tmp_path):
def boom(*a, **k):
raise requests.ConnectionError("refused")
result = codex_model_sync_args({"CODEX_HOME": str(tmp_path)}, "http://localhost:4000", "sk-key", get=boom)
assert isinstance(result, ModelSyncSkipped)
assert "refused" in result.reason
assert not (tmp_path / "litellm-models.json").exists()
@pytest.mark.parametrize(
("response", "reason"),
[(_FakeResponse(500), "HTTP 500"), (_FakeResponse(200, {"data": "nope"}), "unexpected body")],
)
def test_bad_response_is_reported(self, tmp_path, response, reason):
result = codex_model_sync_args(
{"CODEX_HOME": str(tmp_path)}, "http://localhost:4000", "sk-key", get=lambda *a, **k: response
)
assert isinstance(result, ModelSyncSkipped)
assert reason in result.reason
@pytest.mark.parametrize("binary", ["codex", "/opt/bin/codex", "codex.cmd", "/c/npm/codex.CMD"])
def test_codex_syncs_through_the_agent_dispatch_with_the_binary_it_will_run(self, tmp_path, binary):
run = _FakeRun()
result = agent_model_sync_env(
binary,
{"CODEX_HOME": str(tmp_path)},
"http://localhost:4000",
"sk-key",
False,
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
run=run,
)
assert self._catalog_path(result) == str(tmp_path / "litellm-models.json")
assert len(run.calls) == 2
assert all(binary in command for command, _ in run.calls)
def test_opencode_dispatch_never_runs_codex(self):
def boom(*a, **k):
raise AssertionError("only the Codex sync reads its catalog back")
result = agent_model_sync_env(
"opencode",
{},
"http://localhost:4000",
"sk-key",
False,
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
run=boom,
)
assert "OPENCODE_CONFIG_CONTENT" in result
def test_codex_lists_its_own_models_then_reads_the_catalog_back_before_launch(self, tmp_path):
run = _FakeRun()
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=run)
path = self._catalog_path(result)
assert [command for command, _ in run.calls] == [
("codex", "debug", "models"),
("codex", "-c", f"model_catalog_json={json.dumps(path)}", "debug", "models"),
]
for _, options in run.calls:
assert options["env"] == {"CODEX_HOME": str(tmp_path)}
assert options["stdin"] is subprocess.DEVNULL
assert options["capture_output"] is True
assert options["encoding"] == "utf-8"
assert options["timeout"] == 10
def test_codex_rejecting_the_catalog_skips_the_sync_and_keeps_the_file(self, tmp_path):
stderr = (
"Error: failed to parse model_catalog_json path `/home/me/.codex/litellm-models.json` as JSON: "
"missing field `supports_parallel_tool_calls` at line 1 column 21648\n"
)
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=_FakeRun(1, stderr))
assert isinstance(result, ModelSyncSkipped)
assert result.reason == (
"`codex debug models` exited 1: Error: failed to parse model_catalog_json path "
"`/home/me/.codex/litellm-models.json` as JSON: missing field `supports_parallel_tool_calls` "
"at line 1 column 21648"
)
assert (tmp_path / "litellm-models.json").exists()
def test_codex_without_debug_models_skips_the_sync(self, tmp_path):
stderr = "error: unrecognized subcommand 'models'\n\nUsage: codex debug [OPTIONS] <COMMAND>\n"
run = _FakeRun(2, stderr, stock=None)
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=run)
assert isinstance(result, ModelSyncSkipped)
assert result.reason == "`codex debug models` exited 2: error: unrecognized subcommand 'models'"
assert len(run.calls) == 1
assert not (tmp_path / "litellm-models.json").exists()
def test_codex_failing_silently_is_reported(self, tmp_path):
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=_FakeRun(1))
assert isinstance(result, ModelSyncSkipped)
assert result.reason == "`codex debug models` exited 1: no output"
@pytest.mark.parametrize(
"error", [OSError("codex vanished"), subprocess.TimeoutExpired("codex", 10)], ids=["oserror", "timeout"]
)
def test_unrunnable_preflight_is_reported_not_raised(self, tmp_path, error):
def failing_run(*a, **k):
raise error
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=failing_run)
assert isinstance(result, ModelSyncSkipped)
assert result.reason.startswith("`codex debug models` failed: ")
assert str(error) in result.reason
def test_windows_shim_preflight_goes_through_cmd_exe(self, tmp_path):
shim = _WINDOWS_CLAUDE_CMD.replace("claude", "codex")
run = _FakeRun()
result = codex_model_sync_args(
{"CODEX_HOME": str(tmp_path)},
"http://localhost:4000",
"sk-key",
binary=shim,
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
run=run,
)
override = f"model_catalog_json={json.dumps(self._catalog_path(result))}"
doubled = override.replace('"', '""')
assert [command for command, _ in run.calls] == [
f'{_CMD_PREFIX}""{shim}" "debug" "models""',
f'{_CMD_PREFIX}""{shim}" "-c" "{doubled}" "debug" "models""',
]
def test_default_binary_is_codex_on_path(self):
assert _default_of(codex_model_sync_args, "binary") == "codex"
def test_default_runner_is_subprocess_run(self):
assert _default_of(codex_model_sync_args, "run") is subprocess.run
assert _default_of(agent_model_sync_env, "run") is subprocess.run
def test_skip_verify_keeps_the_launch_offline(self):
def boom(*a, **k):
raise AssertionError("--skip-verify must not touch the proxy")
result = agent_model_sync_env("codex", {}, "http://localhost:4000", "sk-key", True, get=boom)
assert isinstance(result, ModelSyncSkipped)
assert "--skip-verify" in result.reason
def test_default_http_client_is_requests_get(self):
assert _default_of(codex_model_sync_args, "get") is requests.get
class TestRunAgent:
def test_synced_args_precede_user_args_and_follow_provider_overrides(self):
calls = {}
run_agent(
"http://localhost:4000",
"sk-key",
["codex", "exec", "hi"],
base_env={},
sync_models=lambda *a: ModelSyncArgs(("-c", 'model_catalog_json="/tmp/c.json"')),
which=lambda name: "/usr/local/bin/codex",
verify=lambda *a: None,
launcher=lambda p, a, e: calls.update(args=tuple(a), env=dict(e)),
)
args = calls["args"]
assert args[-2:] == ("exec", "hi")
assert args[args.index('model_catalog_json="/tmp/c.json"') - 1] == "-c"
assert (
args.index('model_provider="litellm"') < args.index('model_catalog_json="/tmp/c.json"') < args.index("exec")
)
assert calls["env"]["OPENAI_API_KEY"] == "sk-key"
assert "model_catalog_json" not in json.dumps(calls["env"])
def test_synced_model_config_reaches_the_agent_alongside_profile_env(self):
calls = {}
run_agent(
@ -405,7 +863,13 @@ class TestRunAgent:
launcher=lambda p, a, e: order.append("launch"),
)
assert order == ["verify", "sync", "launch"]
assert calls["args"] == ("opencode", {"HOME": "/home/me"}, "http://localhost:4000", "sk-key", False)
assert calls["args"] == (
"/usr/local/bin/opencode",
{"HOME": "/home/me"},
"http://localhost:4000",
"sk-key",
False,
)
def test_unreachable_proxy_is_not_asked_for_models(self):
def failing_verify(*a):
@ -1050,10 +1514,7 @@ class TestAgentCommands:
assert captured["api_key"] == "sk-key"
assert captured["command"] == ["claude", "--resume", "-p", "hi"]
assert captured["skip_verify"] is False
assert (
"routing Claude Code through proxy at http://localhost:4000"
in result.output
)
assert "routing Claude Code through proxy at http://localhost:4000" in result.output
def test_codex_shows_friendly_name(self):
captured = {}
@ -1126,14 +1587,10 @@ class TestAgentCommands:
with (
patch(f"{AGENTS_MODULE}._is_interactive", return_value=True),
patch(f"{AGENTS_MODULE}.login", fake_login),
patch(
f"{AGENTS_MODULE}.get_stored_api_key", return_value="sk-after-login"
) as mock_get,
patch(f"{AGENTS_MODULE}.get_stored_api_key", return_value="sk-after-login") as mock_get,
patch(
f"{AGENTS_MODULE}.run_agent",
side_effect=lambda base_url, api_key, command, **k: captured.update(
api_key=api_key
),
side_effect=lambda base_url, api_key, command, **k: captured.update(api_key=api_key),
),
):
result = self.runner.invoke(

View file

@ -1,5 +1,6 @@
import json
from datetime import datetime, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -266,6 +267,51 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buff
assert popped_keys[6] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
@pytest.mark.asyncio
async def test_org_member_spend_is_summed_across_pods_and_restored_on_rpush_failure(
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock
):
from litellm.proxy._types import Litellm_EntityType
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
DailySpendUpdateQueue,
)
from litellm.proxy.db.db_transaction_queue.spend_update_queue import (
SpendUpdateQueue,
)
member_key: Final = "organization_id::org-1::user_id::user-1"
pod_json: Final = json.dumps({"org_member_list_transactions": {member_key: 0.25}})
mock_redis_cache.async_lpop_pipeline = AsyncMock(
return_value=[[pod_json, pod_json], None, None, None, None, None, None]
)
(db_spend, *_rest) = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
assert db_spend is not None
assert db_spend["org_member_list_transactions"] == {member_key: 0.5}
mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away"))
spend_queue: Final = SpendUpdateQueue()
await spend_queue.add_update(
{
"entity_type": Litellm_EntityType.ORGANIZATION_MEMBER,
"entity_id": member_key,
"response_cost": 1.5,
}
)
await redis_update_buffer.store_in_memory_spend_updates_in_redis(
spend_update_queue=spend_queue,
daily_spend_update_queue=DailySpendUpdateQueue(),
daily_team_spend_update_queue=DailySpendUpdateQueue(),
daily_org_spend_update_queue=DailySpendUpdateQueue(),
daily_end_user_spend_update_queue=DailySpendUpdateQueue(),
daily_agent_spend_update_queue=DailySpendUpdateQueue(),
)
restored_spend: Final = await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
assert restored_spend["org_member_list_transactions"] == {member_key: 1.5}
@pytest.mark.asyncio
async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis():
"""When redis_cache is None, should return all Nones"""

View file

@ -8,6 +8,7 @@ from collections.abc import Callable
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, call, patch
import pytest
@ -944,6 +945,121 @@ async def test_commit_spend_updates_to_db_increments_team_member_spend_and_total
}
@pytest.mark.asyncio
async def test_org_spend_increments_organization_membership_row_for_the_calling_user():
"""A request made with a user_id inside an org must increment that user's
LiteLLM_OrganizationMembership.spend, not only the org total, or the
Organizations > Members UI renders '-' for every member."""
db_writer: Final = DBSpendUpdateWriter()
await db_writer._update_org_db(
response_cost=0.75,
org_id="org-abc",
user_id="user-xyz",
prisma_client=MagicMock(),
)
transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
mock_batcher: Final = MagicMock()
mock_prisma_client: Final = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(mock_batcher))
proxy_logging: Final = MagicMock()
proxy_logging.call_details = {}
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=proxy_logging,
db_spend_update_transactions=transactions,
)
mock_batcher.litellm_organizationtable.update_many.assert_called_once_with(
where={"organization_id": "org-abc"},
data={"spend": {"increment": 0.75}},
)
mock_batcher.litellm_organizationmembership.update_many.assert_called_once_with(
where={"organization_id": "org-abc", "user_id": "user-xyz"},
data={"spend": {"increment": 0.75}},
)
@pytest.mark.asyncio
async def test_org_spend_without_user_id_leaves_organization_membership_untouched():
db_writer: Final = DBSpendUpdateWriter()
await db_writer._update_org_db(
response_cost=0.75,
org_id="org-abc",
user_id=None,
prisma_client=MagicMock(),
)
transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
mock_batcher: Final = MagicMock()
mock_prisma_client: Final = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(mock_batcher))
proxy_logging: Final = MagicMock()
proxy_logging.call_details = {}
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=proxy_logging,
db_spend_update_transactions=transactions,
)
mock_batcher.litellm_organizationtable.update_many.assert_called_once()
mock_batcher.litellm_organizationmembership.update_many.assert_not_called()
@pytest.mark.asyncio
async def test_org_spend_keeps_member_attribution_when_ids_contain_the_key_delimiter():
db_writer: Final = DBSpendUpdateWriter()
await db_writer._update_org_db(
response_cost=0.75,
org_id="division::west",
user_id="user::42",
prisma_client=MagicMock(),
)
transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
mock_batcher: Final = MagicMock()
mock_prisma_client: Final = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(mock_batcher))
proxy_logging: Final = MagicMock()
proxy_logging.call_details = {}
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=proxy_logging,
db_spend_update_transactions=transactions,
)
mock_batcher.litellm_organizationmembership.update_many.assert_called_once_with(
where={"organization_id": "division::west", "user_id": "user::42"},
data={"spend": {"increment": 0.75}},
)
@pytest.mark.asyncio
async def test_batch_database_updates_queues_org_member_spend_for_the_request_user():
db_writer: Final = DBSpendUpdateWriter()
await db_writer._batch_database_updates(
response_cost=0.1,
user_id="u1",
hashed_token="t1",
team_id=None,
org_id="org1",
end_user_id=None,
prisma_client=MagicMock(),
litellm_proxy_budget_name=None,
payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.1},
)
transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
assert transactions["org_list_transactions"] == {"org1": 0.1}
assert transactions["org_member_list_transactions"] == {"organization_id::org1::user_id::u1": 0.1}
@pytest.mark.asyncio
async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id():
"""
@ -2904,6 +3020,7 @@ async def test_update_daily_spend_retries_deadlock(monkeypatch):
("team_list_transactions", "team-1"),
("team_member_list_transactions", "team_id::team-1::user_id::user-1"),
("org_list_transactions", "org-1"),
("org_member_list_transactions", "organization_id::org-1::user_id::user-1"),
("tag_list_transactions", "tag-1"),
("agent_list_transactions", "agent-1"),
],

View file

@ -6311,7 +6311,7 @@ async def test_an_open_circuit_breaker_reads_the_sliding_window_locally_without_
(
{
"team_id": "t",
"metadata": {"model_rpm_limit": {"test-model": 100}},
"metadata": {"model_rpm_limit": {"other-model": 100}},
"team_metadata": {"model_rpm_limit": {"test-model": 1}},
},
{},
@ -6529,3 +6529,113 @@ async def test_request_capacity_rejection_keeps_existing_redis_mirror():
pytest.fail("rejection released another request's mirrored slot")
assert exc.value.status_code == 429
assert await cache.async_get_cache(counter_key, local_only=True) == 1
@pytest.mark.parametrize(
"key_limits",
[
{"metadata": {"model_rpm_limit": {"test-model": 3}}},
{"model_max_budget": {"test-model": {"rpm_limit": 3}}},
],
)
@pytest.mark.asyncio
async def test_key_model_rpm_override_takes_precedence_over_team_model_rpm_limit(key_limits):
cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
auth = UserAPIKeyAuth(
api_key=hash_token("sk-key-override"),
team_id="t",
team_metadata={"model_rpm_limit": {"test-model": 1}},
**key_limits,
)
async def request():
await handler.async_pre_call_hook(
user_api_key_dict=auth, cache=cache, data={"model": "test-model"}, call_type="acompletion"
)
for _ in range(3):
await request()
with pytest.raises(HTTPException) as exc:
await request()
assert exc.value.status_code == 429
assert "model_per_key" in str(exc.value.detail)
@pytest.mark.parametrize(
"key_limits, override_key_gets_through",
[
({"model_rpm_limit": {"test-model": 10}}, False),
({"model_rpm_limit": {"test-model": 10}, "model_tpm_limit": {"test-model": 5000}}, True),
],
ids=["rpm_only_override_still_shares_team_tpm", "rpm_and_tpm_override_leaves_team_tpm"],
)
@pytest.mark.asyncio
async def test_key_model_rpm_override_keeps_team_model_tpm_limit(key_limits, override_key_gets_through):
cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
team_metadata = {"model_rpm_limit": {"test-model": 5}, "model_tpm_limit": {"test-model": 500}}
sibling_key = UserAPIKeyAuth(api_key=hash_token("sk-sibling"), team_id="t", team_metadata=team_metadata)
override_key = UserAPIKeyAuth(
api_key=hash_token("sk-key-override"), team_id="t", metadata=key_limits, team_metadata=team_metadata
)
async def request(auth):
await handler.async_pre_call_hook(
user_api_key_dict=auth,
cache=cache,
data={"model": "test-model", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 300},
call_type="acompletion",
)
await request(sibling_key)
if override_key_gets_through:
await request(override_key)
return
with pytest.raises(HTTPException) as exc:
await request(override_key)
assert exc.value.status_code == 429
assert "model_per_team" in str(exc.value.detail)
assert exc.value.headers["rate_limit_type"] == "tokens"
@pytest.mark.parametrize(
"key_metadata, charges_team_model_pool",
[
({}, True),
({"model_rpm_limit": {"test-model": 10}}, True),
({"model_tpm_limit": {"test-model": 5000}}, False),
({"model_tpm_limit": {"other-model": 5000}}, True),
],
ids=["no_override", "rpm_only_override", "tpm_override", "tpm_override_on_other_model"],
)
def test_success_tpm_accounting_skips_team_model_pool_when_key_owns_model_tpm_limit(
key_metadata, charges_team_model_pool
):
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
response = ModelResponse(
id="team-pool-tpm",
object="chat.completion",
created=int(datetime.now().timestamp()),
model="test-model",
usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
choices=[],
)
kwargs = {
"standard_logging_object": {"metadata": {"user_api_key_hash": hash_token("sk-pool"), "user_api_key_team_id": "t"}},
"litellm_params": {
"metadata": {
"model_group": "test-model",
"user_api_key_metadata": key_metadata,
"user_api_key_team_metadata": {"model_tpm_limit": {"test-model": 500}},
}
},
"model": "test-model",
}
ops = handler._build_success_event_pipeline_operations(kwargs=kwargs, response_obj=response, rate_limit_type="output")
charged_keys = {op["key"] for op in ops}
assert handler.create_rate_limit_keys("model_per_key", f"{hash_token('sk-pool')}:test-model", "tokens") in charged_keys
team_pool_key = handler.create_rate_limit_keys("model_per_team", "t:test-model", "tokens")
assert (team_pool_key in charged_keys) is charges_team_model_pool

View file

@ -36,7 +36,11 @@ from litellm.proxy._types import (
UpdateKeyRequest,
)
from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key
from litellm.proxy.auth.auth_checks import (
_delete_cache_key_object,
_project_cache_key,
jwt_key_mapping_cache_key,
)
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
@ -5132,10 +5136,11 @@ async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch):
class _JWTMappingRow:
def __init__(self, token, jwt_claim_name, jwt_claim_value):
def __init__(self, token, jwt_claim_name, jwt_claim_value, jwt_issuer=None):
self.token = token
self.jwt_claim_name = jwt_claim_name
self.jwt_claim_value = jwt_claim_value
self.jwt_issuer = jwt_issuer
class _CascadingJWTMappingTable:
@ -5226,7 +5231,7 @@ async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypat
),
)
assert recording_evict.cache_keys == ("jwt_key_mapping:email:user@example.com",)
assert recording_evict.cache_keys == (jwt_key_mapping_cache_key("email", "user@example.com", None),)
@pytest.mark.asyncio
@ -13131,11 +13136,11 @@ async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new
_execute_virtual_key_regeneration,
)
stale_cache_key = "jwt_key_mapping:sub:user1"
stale_cache_key = jwt_key_mapping_cache_key("sub", "user1", None)
existing_key = _make_regenerate_existing_key()
mock_prisma_client = _make_regenerate_mock_prisma()
mock_prisma_client.db.litellm_jwtkeymapping.find_many = AsyncMock(
return_value=[MagicMock(jwt_claim_name="sub", jwt_claim_value="user1")]
return_value=[MagicMock(jwt_claim_name="sub", jwt_claim_value="user1", jwt_issuer=None)]
)
mock_prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(
return_value=MagicMock(token="new-hashed-token")

View file

@ -5,7 +5,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.gemini_passthrough_logging_handler import (
GeminiPassthroughLoggingHandler,
@ -397,3 +397,52 @@ class TestGeminiPassthroughLoggingHandler:
assert mock_logging_obj.model_call_details["response_cost"] == expected_cost
assert mock_logging_obj.model_call_details["model"] == "veo-2.0-generate-001"
assert mock_logging_obj.model_call_details["custom_llm_provider"] == "gemini"
def test_interactions_create_response_is_priced_as_gemini(self):
"""Regression for LIT-6896: Gemini API Interactions passthrough must not log zero usage."""
usage = {
"total_tokens": 1030,
"total_input_tokens": 10,
"input_tokens_by_modality": [{"modality": "text", "tokens": 10}],
"total_output_tokens": 1020,
"output_tokens_by_modality": [
{"modality": "text", "tokens": 20},
{"modality": "video", "tokens": 1000},
],
"total_tool_use_tokens": 0,
"total_thought_tokens": 0,
}
mock_httpx_response = MagicMock(spec=httpx.Response)
mock_httpx_response.json.return_value = {
"id": "interactions/abc",
"model": "gemini-omni-flash-preview",
"status": "completed",
"usage": usage,
}
mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj)
mock_logging_obj.model_call_details = {}
mock_logging_obj.litellm_call_id = "call-6896"
result = GeminiPassthroughLoggingHandler.gemini_passthrough_handler(
httpx_response=mock_httpx_response,
response_body=mock_httpx_response.json.return_value,
logging_obj=mock_logging_obj,
url_route="https://generativelanguage.googleapis.com/v1beta/interactions",
result="",
start_time=self.start_time,
end_time=self.end_time,
cache_hit=False,
request_body={"model": "gemini-omni-flash-preview", "input": "make a clip"},
)
model_info = litellm.get_model_info(model="gemini-omni-flash-preview", custom_llm_provider="gemini")
expected_cost = (
10 * model_info["input_cost_per_token"]
+ 20 * model_info["output_cost_per_token"]
+ 1000 * model_info["output_cost_per_video_token"]
)
assert result["result"].id == "call-6896"
assert result["result"].usage.completion_tokens_details.video_tokens == 1000
assert result["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert result["kwargs"]["custom_llm_provider"] == "gemini"
assert mock_logging_obj.model_call_details["custom_llm_provider"] == "gemini"

View file

@ -40,6 +40,7 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
llm_passthrough_factory_proxy_route,
milvus_proxy_route,
mistral_proxy_route,
relay_nvidia_nim_request,
openai_proxy_route,
vertex_discovery_proxy_route,
vertex_proxy_route,
@ -5375,6 +5376,186 @@ class TestRouterModelRelayUpstreamContract:
assert result.headers["x-ms-request-id"] == "req-1"
NIM_INFER_BODY = {
"input": [
{"type": "image_url", "url": "data:image/png;base64,AAAA"},
{"type": "image_url", "url": "data:image/png;base64,BBBB"},
]
}
class TestNvidiaNimProxyRoute:
def _request(self) -> MagicMock:
request = MagicMock(spec=Request)
request.method = "POST"
request.headers = {"content-type": "application/json"}
request.query_params = {}
return request
def _recording_router(self, captured: list[dict], deployments: dict[str, str]):
class RecordingRouter:
def get_model_list(self):
return [{"model_name": name, "litellm_params": {"model": model}} for name, model in deployments.items()]
async def allm_passthrough_route(self, **kwargs):
captured.append(kwargs)
return httpx.Response(
200, json={"data": [{"index": 0, "bounding_boxes": {}}]}, headers={"x-nim-request": "r1"}
)
return RecordingRouter()
async def _relay(self, llm_router, endpoint: str, body: dict, user_api_key_dict=None) -> Response:
return await relay_nvidia_nim_request(
llm_router=llm_router,
endpoint=endpoint,
request=self._request(),
request_body=dict(body),
user_api_key_dict=user_api_key_dict or UserAPIKeyAuth(api_key="hashed-token"),
)
@pytest.mark.asyncio
async def test_model_group_in_the_path_selects_the_deployment_and_the_body_stays_model_free(self):
captured: list[dict] = []
router = self._recording_router(
captured,
{
"nim-page-elements": "nvidia_nim/nvidia/nemoretriever-page-elements-v2",
"nim-table": "nvidia_nim/nvidia/nemoretriever-table-structure-v1",
},
)
result = await self._relay(
router,
"nim-page-elements/v1/infer",
NIM_INFER_BODY,
UserAPIKeyAuth(api_key="hashed-token", team_id="team-1"),
)
(relay,) = captured
assert relay["model"] == "nim-page-elements"
assert relay["endpoint"] == "nim-page-elements/v1/infer"
assert relay["method"] == "POST"
assert relay["json"] == NIM_INFER_BODY
assert "model" not in relay["json"]
assert relay["litellm_metadata"]["user_api_key_team_id"] == "team-1"
assert result.status_code == 200
assert json.loads(result.body) == {"data": [{"index": 0, "bounding_boxes": {}}]}
assert result.headers["x-nim-request"] == "r1"
@pytest.mark.asyncio
async def test_model_group_with_a_slash_is_matched_as_the_longest_leading_path(self):
captured: list[dict] = []
router = self._recording_router(
captured, {"nvidia/nemoretriever-page-elements-v2": "nvidia_nim/nvidia/nemoretriever-page-elements-v2"}
)
await self._relay(router, "nvidia/nemoretriever-page-elements-v2/v1/infer", NIM_INFER_BODY)
assert captured[0]["model"] == "nvidia/nemoretriever-page-elements-v2"
@pytest.mark.asyncio
async def test_custom_llm_provider_marks_a_deployment_as_nim_without_the_model_prefix(self):
captured: list[dict] = []
class ProviderRouter:
def get_model_list(self):
return [
{
"model_name": "page-elements",
"litellm_params": {
"model": "nvidia/nemoretriever-page-elements-v2",
"custom_llm_provider": "nvidia_nim",
},
}
]
async def allm_passthrough_route(self, **kwargs):
captured.append(kwargs)
return httpx.Response(200, json={"data": []})
await self._relay(ProviderRouter(), "page-elements/v1/infer", NIM_INFER_BODY)
assert captured[0]["model"] == "page-elements"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"endpoint",
["v1/infer", "unknown-group/v1/infer", "nim-page-elements-v2/v1/infer", "gpt-4o/v1/infer"],
)
async def test_path_without_a_nim_model_group_is_rejected_before_any_upstream_call(self, endpoint):
captured: list[dict] = []
router = self._recording_router(
captured,
{"nim-page-elements": "nvidia_nim/nvidia/nemoretriever-page-elements-v2", "gpt-4o": "openai/gpt-4o"},
)
with pytest.raises(HTTPException) as exc_info:
await self._relay(router, endpoint, NIM_INFER_BODY)
assert exc_info.value.status_code == 400
assert captured == []
@pytest.mark.asyncio
async def test_a_group_mixing_nim_and_other_deployments_is_rejected_before_any_upstream_call(self):
captured: list[dict] = []
class MixedRouter:
def get_model_list(self):
return [
{
"model_name": "detect",
"litellm_params": {"model": "nvidia_nim/nvidia/nemoretriever-page-elements-v2"},
},
{"model_name": "detect", "litellm_params": {"model": "openai/gpt-4o"}},
]
async def allm_passthrough_route(self, **kwargs):
captured.append(kwargs)
return httpx.Response(200, json={"data": []})
with pytest.raises(HTTPException) as exc_info:
await self._relay(MixedRouter(), "detect/v1/infer", NIM_INFER_BODY)
assert exc_info.value.status_code == 400
assert captured == []
@pytest.mark.asyncio
async def test_no_router_is_rejected_before_any_upstream_call(self):
with pytest.raises(HTTPException) as exc_info:
await self._relay(None, "nim-page-elements/v1/infer", NIM_INFER_BODY)
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_upstream_rejection_is_relayed_with_its_status_body_and_headers(self):
upstream_body = {"detail": "input[0].url must be a data URL"}
class RejectingRouter:
def get_model_list(self):
return [
{
"model_name": "nim-page-elements",
"litellm_params": {"model": "nvidia_nim/nvidia/nemoretriever-page-elements-v2"},
}
]
async def allm_passthrough_route(self, **kwargs):
upstream_request = httpx.Request("POST", "http://nim.internal:8000/v1/infer")
upstream = httpx.Response(
422, json=upstream_body, headers={"x-nim-request": "r2"}, request=upstream_request
)
raise httpx.HTTPStatusError("422", request=upstream_request, response=upstream)
result = await self._relay(
RejectingRouter(), "nim-page-elements/v1/infer", {"input": [{"type": "image_url", "url": "x"}]}
)
assert result.status_code == 422
assert json.loads(result.body) == upstream_body
assert result.headers["x-nim-request"] == "r2"
@pytest.mark.asyncio
async def test_bedrock_count_tokens_error_forwards_provider_headers():
"""The count tokens route converts BedrockError into an HTTPException, and dropping the

View file

@ -34,6 +34,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
)
from litellm.proxy.pass_through_endpoints.success_handler import (
@ -496,6 +497,33 @@ def test_is_vertex_route_ignores_plain_predict_path_segment():
)
def test_interactions_create_routes_are_tracked_for_vertex_and_gemini():
"""
Regression for LIT-6896: Interactions API (gemini-omni) passthrough responses
were never handed to the Vertex/Gemini logging handlers, so SpendLogs rows
landed with zero tokens and zero spend. Only the create URL is billable;
GET/DELETE on an interaction id and non-Google `/interactions` URLs stay generic.
"""
handler = PassThroughEndpointLogging()
vertex_create = "https://aiplatform.googleapis.com/v1beta1/projects/p/locations/global/interactions"
gemini_create = "https://generativelanguage.googleapis.com/v1beta/interactions"
assert handler.is_vertex_route(vertex_create) is True
assert handler.is_vertex_route(f"{vertex_create}/abc123") is False
assert handler.is_vertex_route("https://upstream.example.com/api/interactions") is False
assert handler.is_vertex_route("https://upstream.example.com/locations/eu/interactions") is False
assert (
handler.is_vertex_route(
"https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/interactions"
)
is True
)
assert handler.is_gemini_route(gemini_create, custom_llm_provider="gemini") is True
assert handler.is_gemini_route(f"{gemini_create}/abc123", custom_llm_provider="gemini") is False
assert handler.is_gemini_route(gemini_create, custom_llm_provider=None) is False
@pytest.mark.asyncio
async def test_custom_passthrough_predict_path_logs_via_generic_handler():
"""
@ -5934,6 +5962,32 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key
)
@pytest.mark.parametrize("client_metadata_key", ["litellm_metadata", "metadata"])
def test_passthrough_logs_the_resolved_deployment_model_info_over_the_request_body(client_metadata_key: str):
"""A provider route that resolved a router deployment stashes its model_info on request.state. That
deployment, not a model_info the client put in its own body, is what spend logs and metrics attribute
the call to (LIT-1761: passthrough successes carried model_id="")."""
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent"
mock_request.headers = Headers({})
mock_request.scope = {}
mock_request.state = SimpleNamespace(
**{LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY: {"id": "vertex-gemini-38-flash-dep"}}
)
kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
request=mock_request,
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
passthrough_logging_payload=MagicMock(),
logging_obj=MagicMock(),
_parsed_body={client_metadata_key: {"model_info": {"id": "client-forged-id"}}},
litellm_call_id="lit-1761-call-id",
)
assert kwargs["litellm_params"]["metadata"]["model_info"] == {"id": "vertex-gemini-38-flash-dep"}
@pytest.mark.asyncio
async def test_chat_completion_pass_through_endpoint_answers_an_openai_typed_error_for_an_unknown_model(
monkeypatch: pytest.MonkeyPatch,

View file

@ -1,13 +1,19 @@
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import Request
from starlette.datastructures import Headers, State
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
VertexAIPassThroughHandler,
_base_vertex_proxy_route,
_resolve_vertex_model_from_router,
_upstream_headers_for_vertex_route,
)
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
HttpPassThroughEndpointHelpers,
)
from litellm.types.router import DeploymentTypedDict
@ -758,3 +764,115 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url():
assert (
"gemini-3-pro" in target_url
), f"Actual Vertex AI model name should be in target URL. Got: {target_url}"
@pytest.mark.asyncio
async def test_vertex_passthrough_attributes_the_call_to_the_resolved_deployment():
"""The router deployment that rewrote the upstream URL is the one the logging kwargs must name, so
the Prometheus model_id label (and SpendLogs.model_id) on a Vertex passthrough success reads the
deployment's id instead of "" (LIT-1761)."""
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent"
mock_request.headers = Headers({})
mock_request.scope = {}
mock_request.state = State()
mock_handler = MagicMock()
mock_handler.get_default_base_target_url.return_value = "https://aiplatform.googleapis.com"
mock_router = MagicMock()
mock_router.get_available_deployment_for_pass_through.return_value = {
"model_name": "gemini-3.8-flash",
"litellm_params": {
"model": "vertex_ai/gemini-3.8-flash",
"vertex_project": "p",
"vertex_location": "global",
"use_in_pass_through": True,
},
"model_info": {"id": "vertex-gemini-38-flash-dep"},
}
async def relay_returning_logging_kwargs(
request: Request, fastapi_response: object, user_api_key_dict: UserAPIKeyAuth
) -> dict:
return HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
request=request,
user_api_key_dict=user_api_key_dict,
passthrough_logging_payload=MagicMock(),
logging_obj=MagicMock(),
_parsed_body={"contents": [{"role": "user", "parts": [{"text": "hi"}]}]},
litellm_call_id="lit-1761-call-id",
)
with (
patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it
"litellm.proxy.proxy_server.llm_router", mock_router
),
patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
) as mock_pt_router,
patch( # test-quality-ok: the route offers no injection point for its header preparation
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers",
new_callable=AsyncMock,
return_value=({}, False, "p", "global"),
),
patch( # test-quality-ok: the relay is captured here to read the logging kwargs, the route offers no seam
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
return_value=relay_returning_logging_kwargs,
),
patch( # test-quality-ok: the route calls auth directly rather than through Depends
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
new_callable=AsyncMock,
return_value=UserAPIKeyAuth(api_key="hashed-key"),
),
):
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
logging_kwargs = await _base_vertex_proxy_route(
endpoint="v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
request=mock_request,
fastapi_response=MagicMock(),
get_vertex_pass_through_handler=mock_handler,
)
assert logging_kwargs["litellm_params"]["metadata"]["model_info"]["id"] == "vertex-gemini-38-flash-dep"
def _router_without_deployment() -> MagicMock:
router = MagicMock()
router.get_available_deployment_for_pass_through.return_value = None
return router
def _router_raising_on_lookup() -> MagicMock:
router = MagicMock()
router.get_available_deployment_for_pass_through.side_effect = ValueError("no healthy deployment")
return router
@pytest.mark.parametrize(
"llm_router",
[None, _router_without_deployment(), _router_raising_on_lookup()],
ids=["no-router", "no-matching-deployment", "lookup-raises"],
)
def test_vertex_passthrough_without_a_resolved_deployment_keeps_the_url_and_reports_no_model_info(
llm_router: MagicMock | None,
):
"""A Vertex passthrough call that no router deployment serves must keep the URL-derived values and carry no
deployment model_info, so logging cannot attribute it to a deployment that never handled it."""
resolved = _resolve_vertex_model_from_router(
model_id="gemini-3.8-flash",
llm_router=llm_router,
encoded_endpoint="/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
endpoint="v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
vertex_project="url-project",
vertex_location="url-location",
)
assert resolved == (
"/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
"v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
"url-project",
"url-location",
None,
)

View file

@ -5353,6 +5353,87 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_tokens():
assert all(key not in rows[2] for key in token_keys)
@pytest.mark.asyncio
async def test_build_ui_spend_logs_response_sums_multi_round_session_duration():
"""
Regression test: a multi-round session collapses into a single UI row, so that row
must carry the duration of every round summed, not just the representative call's.
Rows written before request_duration_ms existed are NULL, so the aggregate falls back
to endTime - startTime for them.
"""
from litellm.proxy.spend_tracking.spend_management_endpoints import (
_build_ui_spend_logs_response,
)
session_id = "sess-multi-round-duration"
api_key = "hashed-key-xyz"
dict_rows = [
{
"request_id": "req-1",
"session_id": session_id,
"call_type": "completion",
"api_key": api_key,
"spend": 0.01,
"request_duration_ms": 1200,
},
{
"request_id": "req-2",
"session_id": session_id,
"call_type": "completion",
"api_key": api_key,
"spend": 0.02,
"request_duration_ms": 4200,
},
{
"request_id": "req-3",
"session_id": None,
"call_type": "completion",
"api_key": api_key,
"spend": 0.03,
"request_duration_ms": 900,
},
]
mock_prisma = MagicMock()
mock_prisma.db.query_raw = AsyncMock(
return_value=[
{
"session_id": session_id,
"api_key": api_key,
"session_total_count": 2,
"session_total_spend": 0.03,
"session_total_duration_ms": 5400,
"mcp_tool_call_count": 0,
"mcp_tool_call_spend": 0.0,
}
]
)
result = await _build_ui_spend_logs_response(
prisma_client=mock_prisma,
data=dict_rows,
total_records=3,
page=1,
page_size=50,
total_pages=1,
enrich_session_counts=True,
)
rows = result["data"]
session_rows = rows[:2]
assert [row["session_total_duration_ms"] for row in session_rows] == [5400, 5400]
assert all(isinstance(row["session_total_duration_ms"], int) for row in session_rows)
assert [row["request_duration_ms"] for row in rows] == [1200, 4200, 900]
assert "session_total_duration_ms" not in rows[2]
_, call_args, _ = mock_prisma.db.query_raw.mock_calls[0]
sql = " ".join(call_args[0].split())
assert (
'SUM( COALESCE( request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER ) )'
in sql
)
@pytest.mark.asyncio
async def test_build_ui_spend_logs_response_session_cache_hit_count():
"""

View file

@ -13,7 +13,11 @@ from fastapi.responses import JSONResponse, StreamingResponse
import litellm
from litellm._uuid import uuid
from litellm.constants import MAX_LITELLM_CALL_ID_LENGTH, RETURN_RAW_MODEL_NAME_METADATA_KEY
from litellm.constants import (
CLIENT_REQUESTED_MODEL_SCOPE_KEY,
MAX_LITELLM_CALL_ID_LENGTH,
RETURN_RAW_MODEL_NAME_METADATA_KEY,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.opentelemetry import UserAPIKeyAuth
from litellm.proxy.common_request_processing import (
@ -4395,6 +4399,53 @@ class TestDisconnectGatherCleanup:
)
@pytest.mark.asyncio
@pytest.mark.parametrize("client_model, expected", [("AgentX-LLM", "AgentX-LLM"), (None, "gpt-mini")])
async def test_response_model_echoes_the_name_the_client_sent_before_auth_rewrote_it(
monkeypatch, client_model, expected
):
"""LIT-3054: auth resolves router_settings.model_group_alias in the body, so the alias the
client sent only survives in the request scope. The response must still echo it."""
import litellm.proxy.common_request_processing as cpr
async def llm():
return litellm.ModelResponse(
model="gpt-4o-mini", choices=[{"message": {"role": "assistant", "content": "pong"}}]
)
async def fake_route_request(**_kwargs):
return llm()
logging_obj = MagicMock(litellm_call_id="call-id", _defer_async_logging=False)
proxy_logging = MagicMock(spec=ProxyLogging)
proxy_logging.during_call_hook = AsyncMock(return_value=None)
proxy_logging.post_call_success_hook = AsyncMock(side_effect=lambda data, user_api_key_dict, response: response)
proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={})
proxy_logging._callback_capabilities_cache = {}
monkeypatch.setattr(cpr, "route_request", fake_route_request)
processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-mini", "messages": []})
monkeypatch.setattr(
processor, "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gpt-mini"}, logging_obj))
)
monkeypatch.setattr(processor, "_has_post_call_guardrails", MagicMock(return_value=False))
scope = {"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": [], "query_string": b""}
request = Request({**scope, CLIENT_REQUESTED_MODEL_SCOPE_KEY: client_model} if client_model else scope)
response = await processor.base_process_llm_request(
request=request,
fastapi_response=Response(),
user_api_key_dict=ProxyUserAPIKeyAuth(),
proxy_logging_obj=proxy_logging,
general_settings={},
proxy_config=MagicMock(spec=ProxyConfig),
route_type="acompletion",
version=None,
)
assert response.model == expected
class TestStreamingClientDisconnectLogging:
@pytest.mark.asyncio
async def test_record_streaming_client_disconnect_sets_error_information(self):
@ -8169,7 +8220,7 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_
try:
raise exc
except Exception as raised:
_log_llm_api_exception(raised)
_log_llm_api_exception(raised, "call-id-for-traceback-test")
finally:
verbose_proxy_logger.propagate = False
@ -8663,3 +8714,84 @@ class TestBackgroundResponseRetrievalGovernance:
assert "_guardrail_pipelines" not in data["litellm_metadata"]
assert "applied_policies" not in data["litellm_metadata"]
class TestErrorLogCarriesCallId:
"""Regression for LIT-5856 / #37532: the ERROR line emitted for a failed LLM
request must carry the litellm_call_id the client got back in the
x-litellm-call-id response header, so a logged exception can be tied to a
specific request."""
async def _invoke(self, data: dict[str, object]) -> None:
from litellm._logging import verbose_proxy_logger
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
proxy_logging_obj: Final = MagicMock()
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
verbose_proxy_logger.propagate = True
try:
with pytest.raises(ProxyException):
await processor._handle_llm_api_exception(
e=ValueError("upstream blew up"),
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"),
proxy_logging_obj=proxy_logging_obj,
)
finally:
verbose_proxy_logger.propagate = False
@staticmethod
def _error_record(caplog: pytest.LogCaptureFixture):
return next(r for r in caplog.records if "_handle_llm_api_exception(): Exception occured" in r.getMessage())
async def test_call_id_from_logging_obj_is_logged(self, caplog: pytest.LogCaptureFixture) -> None:
call_id: Final = str(uuid.uuid4())
logging_obj: Final = MagicMock()
logging_obj.litellm_call_id = call_id
with caplog.at_level("ERROR", logger="LiteLLM Proxy"):
await self._invoke({"litellm_logging_obj": logging_obj, "litellm_call_id": "stale-id"})
record: Final = self._error_record(caplog)
assert record.litellm_call_id == call_id
assert call_id in record.getMessage()
async def test_call_id_falls_back_to_request_data(self, caplog: pytest.LogCaptureFixture) -> None:
call_id: Final = str(uuid.uuid4())
with caplog.at_level("ERROR", logger="LiteLLM Proxy"):
await self._invoke({"litellm_call_id": call_id})
record: Final = self._error_record(caplog)
assert record.litellm_call_id == call_id
assert call_id in record.getMessage()
async def test_call_id_falls_back_when_logging_obj_has_none(self, caplog: pytest.LogCaptureFixture) -> None:
call_id: Final = str(uuid.uuid4())
logging_obj: Final = MagicMock()
logging_obj.litellm_call_id = None
with caplog.at_level("ERROR", logger="LiteLLM Proxy"):
await self._invoke({"litellm_logging_obj": logging_obj, "litellm_call_id": call_id})
record: Final = self._error_record(caplog)
assert record.litellm_call_id == call_id
assert call_id in record.getMessage()
def test_client_disconnect_log_carries_call_id(self, caplog: pytest.LogCaptureFixture) -> None:
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_request_processing import (
_CLIENT_DISCONNECT_DETAIL,
_log_llm_api_exception,
)
call_id: Final = str(uuid.uuid4())
verbose_proxy_logger.propagate = True
try:
with caplog.at_level("INFO", logger="LiteLLM Proxy"):
_log_llm_api_exception(
HTTPException(status_code=499, detail=_CLIENT_DISCONNECT_DETAIL),
call_id,
)
finally:
verbose_proxy_logger.propagate = False
record: Final = caplog.records[-1]
assert record.litellm_call_id == call_id
assert call_id in record.getMessage()

View file

@ -139,6 +139,35 @@ class TestProxyInitializationHelpers:
)
assert args["timeout_worker_healthcheck"] == 15
@staticmethod
def _uvicorn_access_info_enabled(args: dict) -> bool:
import logging
loggers = tuple(logging.getLogger(n) for n in ("uvicorn", "uvicorn.error", "uvicorn.access", "uvicorn.asgi"))
saved = tuple((lg, lg.handlers[:], lg.level, lg.propagate) for lg in loggers)
try:
uvicorn.Config(**args).configure_logging()
return logging.getLogger("uvicorn.access").isEnabledFor(logging.INFO)
finally:
for lg, handlers, level, propagate in saved:
lg.handlers[:] = handlers
lg.setLevel(level)
lg.propagate = propagate
def test_litellm_log_error_silences_uvicorn_info_lines(self, monkeypatch):
monkeypatch.setenv("LITELLM_LOG", "ERROR")
args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000)
assert "log_config" not in args
assert self._uvicorn_access_info_enabled(args) is False
def test_unset_litellm_log_keeps_uvicorn_default_info_lines(self, monkeypatch):
monkeypatch.delenv("LITELLM_LOG", raising=False)
args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000)
assert "log_level" not in args
assert self._uvicorn_access_info_enabled(args) is True
def test_installed_uvicorn_supports_worker_flags(self):
params = inspect.signature(uvicorn.Config.__init__).parameters
assert "timeout_worker_healthcheck" in params

View file

@ -7401,6 +7401,88 @@ async def test_update_general_settings_apply_user_budget_to_team_keys_yaml_wins(
assert ps.general_settings["apply_user_budget_to_team_keys"] is True
@pytest.mark.asyncio
async def test_update_general_settings_keeps_yaml_pass_through_endpoints_next_to_db_ones():
"""user_api_key_auth honours ``auth: false`` only for entries it finds in
general_settings["pass_through_endpoints"]. The DB overlay used to replace that
list wholesale, so once one endpoint existed in the DB the YAML-declared
auth-disabled route started answering 401 while staying registered."""
from litellm.proxy._types import ProxyException
from litellm.proxy.proxy_server import ProxyConfig
yaml_endpoint: Final = {"path": "/v1/cuopt/request", "target": "https://example.com/post", "auth": False}
db_endpoint: Final = {"id": "db-1", "path": "/v1/db-echo", "target": "https://example.com/post", "auth": True}
def request_without_key(path: str) -> MagicMock:
request: Final = MagicMock()
request.url.path = path
request.headers = {}
request.query_params = {}
return request
settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam
yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]) # test-quality-ok: module global holding the YAML endpoints the fix merges in
initialize: Final = patch("litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here
master_key: Final = patch("litellm.proxy.proxy_server.master_key", "sk-master") # test-quality-ok: a set master key is what makes a missing Authorization header a 401
with settings, yaml_endpoints, initialize, master_key:
await ProxyConfig()._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]})
anonymous: Final = await user_api_key_auth(request=request_without_key("/v1/cuopt/request"), api_key=None)
assert anonymous.api_key is None
with pytest.raises(ProxyException) as still_protected:
await user_api_key_auth(request=request_without_key("/v1/db-echo"), api_key=None)
assert still_protected.value.code == "401"
@pytest.mark.asyncio
@pytest.mark.parametrize(
("db_methods", "yaml_methods"),
[(None, None), (["POST"], ["GET"])],
ids=["all-methods", "disjoint-methods"],
)
async def test_update_general_settings_db_pass_through_endpoint_overrides_yaml_entry_on_the_same_path(
db_methods: list[str] | None, yaml_methods: list[str] | None
):
"""The auth check matches pass-through entries by path only and lets any
matching ``auth: false`` entry through, so a DB ``auth: true`` entry can only
lock down a YAML-declared path if the YAML entry is dropped from the merged
list, whatever ``methods`` either entry declares."""
from litellm.proxy._types import ProxyException
from litellm.proxy.proxy_server import ProxyConfig
yaml_endpoint: Final = {
"path": "/v1/cuopt/request",
"target": "https://example.com/post",
"auth": False,
"methods": yaml_methods,
}
db_endpoint: Final = {
"id": "db-1",
"path": "/v1/cuopt/request",
"target": "https://example.com/post",
"auth": True,
"methods": db_methods,
}
request: Final = MagicMock()
request.url.path = "/v1/cuopt/request"
request.method = "POST"
request.headers = {}
request.query_params = {}
settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam
yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]) # test-quality-ok: module global holding the YAML endpoints the fix merges in
initialize: Final = patch("litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here
master_key: Final = patch("litellm.proxy.proxy_server.master_key", "sk-master") # test-quality-ok: a set master key is what makes a missing Authorization header a 401
with settings, yaml_endpoints, initialize, master_key:
await ProxyConfig()._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]})
with pytest.raises(ProxyException) as locked_down:
await user_api_key_auth(request=request, api_key=None)
assert locked_down.value.code == "401"
def _fill_user_api_key_cache(cache: DualCache, count: int) -> None:
for index in range(count):
cache.set_cache(key=f"key-{index}", value={"token": f"key-{index}"}, local_only=True)
@ -10750,6 +10832,7 @@ def test_get_config_list_includes_anthropic_prompt_caching_fields(monkeypatch):
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
monkeypatch.setattr(litellm, "anthropic_prompt_caching_ttl", "1h")
monkeypatch.setattr(litellm, "openai_system_messages_first", False)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
)
@ -10771,6 +10854,10 @@ def test_get_config_list_includes_anthropic_prompt_caching_fields(monkeypatch):
assert fields["enable_anthropic_prompt_caching"]["field_tab"] == "prompt_caching"
assert fields["anthropic_prompt_caching_ttl"]["field_tab"] == "prompt_caching"
assert fields["budget_exceeded_throttle_percentage"]["field_tab"] is None
assert fields["openai_system_messages_first"]["field_type"] == "Boolean"
assert fields["openai_system_messages_first"]["field_value"] is False
assert fields["openai_system_messages_first"]["field_tab"] == "prompt_caching"
finally:
app.dependency_overrides.clear()
@ -10887,6 +10974,7 @@ def test_general_settings_ui_defaults_unchanged_for_existing_fields():
[
("enable_anthropic_prompt_caching", True),
("anthropic_prompt_caching_ttl", "1h"),
("openai_system_messages_first", True),
],
)
def test_prompt_caching_settings_propagate_on_config_reload(monkeypatch, field_name, db_value):
@ -10945,6 +11033,8 @@ def test_get_config_list_marks_untouched_prompt_caching_flag_as_not_set(monkeypa
("enable_anthropic_prompt_caching", False),
("anthropic_prompt_caching_ttl", "5m"),
("anthropic_prompt_caching_ttl", "1h"),
("openai_system_messages_first", True),
("openai_system_messages_first", False),
],
)
@pytest.mark.asyncio
@ -10993,6 +11083,8 @@ async def test_update_config_field_prompt_caching_persists_to_litellm_settings(m
("anthropic_prompt_caching_ttl", "10m"),
("anthropic_prompt_caching_ttl", "1H"),
("anthropic_prompt_caching_ttl", 3600),
("openai_system_messages_first", "yes"),
("openai_system_messages_first", 1),
],
)
@pytest.mark.asyncio
@ -11032,6 +11124,7 @@ async def test_update_config_field_prompt_caching_rejects_invalid(monkeypatch, f
[
("enable_anthropic_prompt_caching", False),
("anthropic_prompt_caching_ttl", None),
("openai_system_messages_first", False),
("budget_exceeded_throttle_percentage", None),
],
)

View file

@ -1,3 +1,4 @@
import time
from typing import Final
import pytest
@ -7,6 +8,7 @@ import litellm
from litellm.cost_calculator import (
BaseTokenUsageProcessor,
RealtimeAPITokenUsageProcessor,
ResponsesWebSocketTokenUsageProcessor,
completion_cost,
cost_per_token,
handle_realtime_stream_cost_calculation,
@ -15,9 +17,11 @@ from litellm.cost_calculator import (
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo
from litellm.types.llms.base import CachedTokensDetails
from litellm.types.llms.openai import OpenAIRealtimeStreamList
from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponseAPIUsage, ResponsesAPIResponse
from litellm.types.rerank import RerankResponse
from litellm.types.utils import (
CallTypes,
LiteLLMRealtimeStreamLoggingObject,
ModelInfo,
ModelResponse,
PromptTokensDetailsWrapper,
@ -3692,6 +3696,31 @@ def test_completion_cost_bills_interactions_video_output_at_video_rate():
assert cost == pytest.approx(expected)
@pytest.mark.parametrize("video_count", [2, 3])
def test_completion_cost_multiplies_video_cost_by_generated_video_count(video_count: int) -> None:
"""Regression for LIT-6896: a Veo request for N samples generates N videos and must be billed N times."""
from litellm.types.videos.main import VideoObject
def _video(usage: dict[str, object]) -> VideoObject:
return VideoObject(id="v", object="video", status="processing", model="veo-3.1-fast-generate-001", usage=usage)
single_cost = completion_cost(
completion_response=_video({"duration_seconds": 4.0, "video_resolution": "720p"}),
model="veo-3.1-fast-generate-001",
custom_llm_provider="vertex_ai",
call_type="create_video",
)
multi_cost = completion_cost(
completion_response=_video({"duration_seconds": 4.0, "video_resolution": "720p", "video_count": video_count}),
model="veo-3.1-fast-generate-001",
custom_llm_provider="vertex_ai",
call_type="create_video",
)
assert single_cost > 0
assert multi_cost == pytest.approx(single_cost * video_count)
@pytest.mark.parametrize(
"batch_rate,expected_prompt,expected_completion",
[
@ -4718,3 +4747,74 @@ def test_completion_cost_ocr_ignores_deployment_pricing_without_custom_pricing_f
litellm_logging_obj=logging_obj,
)
assert cost == 0.0
def test_completion_cost_prices_responses_websocket_turns_per_service_tier():
"""Issue #41299: a session mixing default and priority turns must price each turn at
its own returned service_tier, not the summed usage at a single tier."""
events = [
{"type": "response.created", "response": {}},
{
"type": "response.completed",
"response": {
"service_tier": "default",
"usage": {"input_tokens": 100, "output_tokens": 40, "total_tokens": 140},
},
},
{"type": "rate_limits.updated", "rate_limits": {}},
{
"type": "response.completed",
"response": {
"service_tier": "priority",
"usage": {"input_tokens": 60, "output_tokens": 10, "total_tokens": 70},
},
},
{"type": "response.failed", "response": {"usage": None}},
]
partition = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(events)
assert tuple(partition.keys()) == ("default", "priority")
assert len(partition["default"]) == 1
assert len(partition["priority"]) == 1
logging_obj = Logging(
model="gpt-5.4",
messages=[],
stream=False,
call_type=CallTypes.aresponses_websocket.value,
start_time=time.time(),
litellm_call_id="responses-ws-tier-test",
function_id="responses-ws-tier-test",
)
normalized = logging_obj.normalize_logging_result(result=events)
assert isinstance(normalized, LiteLLMRealtimeStreamLoggingObject)
assert normalized.service_tier is None
def _http_cost(input_tokens: int, output_tokens: int, service_tier: str) -> float:
return completion_cost(
completion_response=ResponsesAPIResponse(
id=f"resp-{service_tier}",
created_at=1700000000,
output=[],
service_tier=service_tier,
usage=ResponseAPIUsage(
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=input_tokens + output_tokens,
),
),
model="gpt-5.4",
call_type=CallTypes.aresponses.value,
custom_llm_provider="openai",
)
ws_cost = completion_cost(
completion_response=normalized,
model="gpt-5.4",
call_type=CallTypes.aresponses_websocket.value,
custom_llm_provider="openai",
)
assert ws_cost == pytest.approx(_http_cost(100, 40, "default") + _http_cost(60, 10, "priority"))
assert ws_cost != pytest.approx(_http_cost(160, 50, "default"))
assert ws_cost != pytest.approx(_http_cost(160, 50, "priority"))

View file

@ -12206,6 +12206,83 @@ def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_o
@pytest.mark.parametrize(
"model,provider,expected",
[
("anthropic/claude-opus-5", None, True),
("claude-opus-4-8", None, True),
("anthropic/claude-opus-4-7", None, False),
("anthropic/claude-opus-4-6", None, False),
("anthropic/claude-sonnet-5", None, False),
("anthropic/off-map-opus", None, False),
("vertex_ai/claude-opus-5", None, False),
("bedrock/claude-opus-5", None, False),
("claude-opus-5", "vertex_ai", False),
("claude-opus-5", "bedrock", False),
],
)
@pytest.mark.parametrize("operator_flag", [True, False])
def test_model_group_info_fast_mode_uses_exact_provider_catalog(
local_model_cost_map: None, model: str, provider: str | None, expected: bool, operator_flag: bool
) -> None:
router: Final = Router(model_list=[{
"model_name": "fast-group",
"litellm_params": {"model": model, "custom_llm_provider": provider, "api_key": "fake-key"},
"model_info": {"supports_fast_mode": operator_flag},
}])
result: Final = router.get_model_group_info("fast-group")
assert result is not None
assert result.supports_fast_mode is expected
@pytest.mark.parametrize("flag", [None, False, "true", 1])
def test_model_group_info_fast_mode_fails_closed_without_explicit_boolean(
local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, flag: object
) -> None:
entry: Final = {key: value for key, value in litellm.model_cost["claude-opus-5"].items()
if key != "supports_fast_mode"}
if flag is not None:
entry["supports_fast_mode"] = flag
monkeypatch.setitem(litellm.model_cost, "claude-opus-5", entry)
router: Final = Router(model_list=[{
"model_name": "fast-group",
"litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "fake-key"},
"model_info": {"supports_fast_mode": True},
}])
result: Final = router.get_model_group_info("fast-group")
assert result is not None
assert result.supports_fast_mode is False
@pytest.mark.parametrize("other_model,expected", [
("anthropic/claude-opus-4-8", True),
("anthropic/claude-opus-4-7", False),
("anthropic/off-map-opus", False),
("vertex_ai/claude-opus-5", False),
("bedrock/claude-opus-5", False),
])
@pytest.mark.parametrize("reverse", [True, False])
def test_model_group_info_fast_mode_requires_every_deployment(
local_model_cost_map: None, other_model: str, expected: bool, reverse: bool
) -> None:
models: Final = (other_model, "anthropic/claude-opus-5") if reverse else (
"anthropic/claude-opus-5", other_model
)
router: Final = Router(model_list=[{
"model_name": "fast-group",
"litellm_params": {"model": model, "api_key": "fake-key"},
} for model in models])
result: Final = router.get_model_group_info("fast-group")
assert result is not None
assert result.supports_fast_mode is expected
def test_model_group_info_surfaces_supports_parallel_function_calling(local_model_cost_map):
"""``/model_group/info`` folds each deployment's registry flags into the group; a deployment whose
registry entry declares parallel function calling must flip the group to True instead of False."""

View file

@ -902,6 +902,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"supports_sampling_params": {"type": "boolean"},
"supports_output_config": {"type": "boolean"},
"supports_speed": {"type": "boolean"},
"supports_fast_mode": {"type": "boolean"},
"supported_audio_formats": {
"type": "array",
"items": {

View file

@ -71,6 +71,7 @@ const renderWith = (results: DailyData[], overrides: Partial<DailyActivityRange>
isFetchingMore: false,
progress: { currentPage: 1, totalPages: 1 },
cancelled: false,
failed: false,
cancel: vi.fn(),
...overrides,
}}

View file

@ -87,6 +87,7 @@ const CostOptimizationView: React.FC<CostOptimizationViewProps> = ({ accessToken
<PaginationStatusAlerts
isFetchingMore={activity.isFetchingMore}
cancelled={activity.cancelled}
failed={activity.failed}
progress={activity.progress}
cancel={activity.cancel}
/>

View file

@ -35,6 +35,7 @@ describe("PromptCachingTab", () => {
isFetchingMore: false,
progress: { currentPage: 1, totalPages: 1 },
cancelled: false,
failed: false,
cancel: vi.fn(),
};
render(<PromptCachingTab accessToken="test-token" activity={activity} />);

View file

@ -123,6 +123,7 @@ const renderWith = (results: DailyData[], options: RenderOptions = {}) => {
isFetchingMore: false,
progress: { currentPage: 1, totalPages: 1 },
cancelled: false,
failed: false,
cancel: vi.fn(),
}}
/>,

View file

@ -20,6 +20,7 @@ export interface DailyActivityRange {
isFetchingMore: boolean;
progress: { currentPage: number; totalPages: number };
cancelled: boolean;
failed: boolean;
cancel: () => void;
}
@ -64,7 +65,7 @@ export const useScopedDailyActivityRange = (
args: [accessToken, startTime, endTime, userId, true, apiKey],
enabled: !!accessToken && !!startTime && !!endTime,
};
const { data, loading, isFetchingMore, progress, cancelled, cancel } =
const { data, loading, isFetchingMore, progress, cancelled, failed, cancel } =
usePaginatedDailyActivity(activityQueryOptions);
return {
@ -75,6 +76,7 @@ export const useScopedDailyActivityRange = (
isFetchingMore,
progress,
cancelled,
failed,
cancel,
};
};

View file

@ -1,8 +1,10 @@
import * as useAuthorizedModule from "@/app/(dashboard)/hooks/useAuthorized";
import { fireEvent, render, screen, waitFor, within } from "@testing-library/react";
import { fireEvent, screen, waitFor, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { beforeEach, describe, expect, it, vi } from "vitest";
import type { OnUrlUpdateFunction } from "nuqs/adapters/testing";
import { beforeEach, describe, expect, it, Mock, vi } from "vitest";
import { renderWithProviders } from "../../../../../tests/test-utils";
import AllModelsTab from "./AllModelsTab";
import { STATUS_COLUMN_ID, toServerSortField } from "./ModelsTableColumns";
@ -111,6 +113,9 @@ const setModelsInfo = (rows: Record<string, unknown>[], totalCount = rows.length
const lastModelsInfoCall = (): ModelsInfoArgs => modelsInfoCalls[modelsInfoCalls.length - 1];
const lastUrlParams = (onUrlUpdate: Mock<OnUrlUpdateFunction>): URLSearchParams | undefined =>
onUrlUpdate.mock.calls.at(-1)?.[0].searchParams;
const SEARCH_SETTLE_MS = 400;
const MOCK_AUTHORIZED = {
@ -121,6 +126,8 @@ const MOCK_AUTHORIZED = {
userId: "user-123",
userEmail: "test@example.com",
userRole: "Admin",
userRoleLabel: "Admin",
isViewOnly: false,
premiumUser: true,
disabledPersonalKeyCreation: false,
showSSOBanner: false,
@ -149,14 +156,14 @@ describe("AllModelsTab", () => {
it("renders the fetched models and the server row count", async () => {
setModelsInfo([makeRow()], 137);
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(await screen.findByText("gpt-4")).toBeInTheDocument();
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-50 of 137");
});
it("does not re-query after the mount-time debounced search settles unchanged", async () => {
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
const callsAfterMount = modelsInfoCalls.length;
await new Promise((resolve) => setTimeout(resolve, SEARCH_SETTLE_MS));
@ -166,14 +173,14 @@ describe("AllModelsTab", () => {
it("shows the empty state when the proxy returns no models", () => {
setModelsInfo([], 0);
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(screen.getByText("No models found")).toBeInTheDocument();
});
it("shows the loading skeleton while the first page is in flight", () => {
setModelsInfo([], 0, true);
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0);
expect(screen.queryByText("No models found")).not.toBeInTheDocument();
@ -197,7 +204,7 @@ describe("AllModelsTab", () => {
it.each(cases)("sorts %s using the server field %s", async (_label, columnId, serverField, firstDirection) => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(sortHeader(columnId));
await expectIndicator(columnId, firstDirection);
@ -212,7 +219,7 @@ describe("AllModelsTab", () => {
it("cycles a sorted column back to unsorted", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(sortHeader("model_info_updated_at"));
await expectIndicator("model_info_updated_at", "asc");
@ -230,7 +237,7 @@ describe("AllModelsTab", () => {
it("queries the selected team and resets to the first page", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(lastModelsInfoCall().teamId).toBeUndefined();
@ -244,8 +251,7 @@ describe("AllModelsTab", () => {
});
it("debounces the model name search into the server query", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "claude" } });
@ -254,9 +260,138 @@ describe("AllModelsTab", () => {
});
});
describe("URL persistence", () => {
it("writes the typed search to the URL and drops the page so a reload keeps the search", async () => {
setModelsInfo([makeRow()], 200);
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { page: "3" }, onUrlUpdate });
expect(lastModelsInfoCall().page).toBe(3);
fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "claude" } });
await waitFor(() => {
expect(lastUrlParams(onUrlUpdate)?.get("model_search")).toBe("claude");
});
expect(lastUrlParams(onUrlUpdate)?.get("page")).toBeNull();
await waitFor(() => {
expect(lastModelsInfoCall().page).toBe(1);
});
});
it("restores the search box and server query from ?model_search= on mount", () => {
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { model_search: "haiku" } });
expect(screen.getByTestId("datatable-search")).toHaveValue("haiku");
expect(lastModelsInfoCall().search).toBe("haiku");
});
it("restores team, sort, page and page size from the URL into the server query", () => {
setModelsInfo([makeRow()], 200);
renderWithProviders(<AllModelsTab {...defaultProps} />, {
searchParams: {
filter_team: "team-1",
sort_by: "model_info_updated_at",
sort_order: "desc",
page: "2",
page_size: "25",
},
});
const expectedQuery: ModelsInfoArgs = {
teamId: "team-1",
sortBy: "updated_at",
sortOrder: "desc",
page: 2,
size: 25,
};
expect(lastModelsInfoCall()).toMatchObject(expectedQuery);
expect(screen.getByTestId("models-team-select")).toHaveTextContent("Engineering");
});
it("restores the access group and view mode from the URL", () => {
renderWithProviders(<AllModelsTab {...defaultProps} />, {
searchParams: { access_group: "sales-team", view_mode: "all" },
});
expect(lastModelsInfoCall().accessGroup).toBe("sales-team");
expect(screen.queryByText(/To access these models/)).not.toBeInTheDocument();
});
it("clamps a hand-edited page and page size into the range the table supports", () => {
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { page: "0", page_size: "5000" } });
expect(lastModelsInfoCall().page).toBe(1);
expect(lastModelsInfoCall().size).toBe(100);
});
it("keeps the default page size when the URL value is not a number", () => {
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { page_size: "lots" } });
expect(lastModelsInfoCall().size).toBe(50);
});
it("ignores a sort_by the table cannot sort by instead of forwarding it to the server", () => {
renderWithProviders(<AllModelsTab {...defaultProps} />, {
searchParams: { sort_by: "litellm_credential_name", sort_order: "desc" },
});
expect(lastModelsInfoCall().sortBy).toBeUndefined();
expect(lastModelsInfoCall().sortOrder).toBeUndefined();
});
it("writes sort changes to the URL with the page cleared", async () => {
setModelsInfo([makeRow()], 200);
const user = userEvent.setup();
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { page: "2" }, onUrlUpdate });
await user.click(screen.getByTestId("sort-header-model_info_updated_at"));
await waitFor(() => {
expect(lastUrlParams(onUrlUpdate)?.get("sort_by")).toBe("model_info_updated_at");
});
expect(lastUrlParams(onUrlUpdate)?.get("sort_order")).toBeNull();
expect(lastUrlParams(onUrlUpdate)?.get("page")).toBeNull();
await user.click(screen.getByTestId("sort-header-model_info_updated_at"));
await waitFor(() => {
expect(lastUrlParams(onUrlUpdate)?.get("sort_order")).toBe("desc");
});
});
it("clears every table param from the URL on drawer reset", async () => {
setModelsInfo([makeRow()], 200);
const user = userEvent.setup();
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<AllModelsTab {...defaultProps} />, {
searchParams: {
model_search: "haiku",
filter_team: "team-1",
sort_by: "model_name",
page: "2",
view_mode: "all",
},
onUrlUpdate,
});
await user.click(screen.getByTestId("datatable-filters-trigger"));
await user.click(await screen.findByTestId("filter-drawer-reset"));
await waitFor(() => {
expect(lastUrlParams(onUrlUpdate)?.toString()).toBe("");
});
expect(screen.getByTestId("datatable-search")).toHaveValue("");
const defaultQuery: ModelsInfoArgs = { search: undefined, teamId: undefined, sortBy: undefined, page: 1 };
await waitFor(() => {
expect(lastModelsInfoCall()).toMatchObject(defaultQuery);
});
});
});
it("applies a public model name filter through the drawer", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(screen.getByTestId("datatable-filters-trigger"));
await user.click(await screen.findByPlaceholderText("Filter by Public Model Name"));
@ -270,7 +405,7 @@ describe("AllModelsTab", () => {
it("renders every row the server returned for the selected model group so rows match the footer total", () => {
setModelsInfo([makeRow(), { ...makeRow({ model_info: { id: "model-2" } }), model_name: "claude-opus" }], 2);
render(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
const table = screen.getByRole("table");
expect(within(table).getByText("claude-opus")).toBeInTheDocument();
@ -280,7 +415,7 @@ describe("AllModelsTab", () => {
it("asks the server for wildcard deployments instead of hiding rows client-side", () => {
setModelsInfo([makeRow(), { ...makeRow({ model_info: { id: "model-2" } }), model_name: "openai/*" }], 2);
render(<AllModelsTab {...defaultProps} selectedModelGroup="wildcard" />);
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="wildcard" />);
expect(lastModelsInfoCall().wildcardOnly).toBe(true);
expect(within(screen.getByRole("table")).getByText("gpt-4")).toBeInTheDocument();
@ -289,7 +424,7 @@ describe("AllModelsTab", () => {
it("asks the server for the selected access group instead of hiding rows client-side", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(lastModelsInfoCall().wildcardOnly).toBe(false);
await user.click(screen.getByTestId("datatable-filters-trigger"));
@ -303,20 +438,20 @@ describe("AllModelsTab", () => {
});
it("asks the server for the exact selected model group so deployments beyond the first page are found", () => {
render(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
expect(lastModelsInfoCall().modelName).toBe("claude-opus");
expect(lastModelsInfoCall().search).toBeUndefined();
});
it.each(["all", "wildcard"])("sends no exact model name for the %s pseudo group", (group) => {
render(<AllModelsTab {...defaultProps} selectedModelGroup={group} />);
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup={group} />);
expect(lastModelsInfoCall().modelName).toBeUndefined();
});
it("keeps the exact model group alongside a typed search", async () => {
render(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
fireEvent.change(screen.getByPlaceholderText("Search model names…"), { target: { value: "opus" } });
@ -326,7 +461,7 @@ describe("AllModelsTab", () => {
it("resets search, filters, team and sorting from the drawer reset button", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} selectedModelGroup="gpt-4" />);
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="gpt-4" />);
await user.click(screen.getByTestId("models-team-select"));
await user.click(await screen.findByRole("option", { name: "Engineering" }));
@ -343,7 +478,7 @@ describe("AllModelsTab", () => {
it("opens the delete modal from the row and deletes the model", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(await screen.findByTestId("model-delete-model-1"));
expect(await screen.findByText("Delete Model")).toBeInTheDocument();
@ -357,7 +492,7 @@ describe("AllModelsTab", () => {
it("pauses a model through the row toggle", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(await screen.findByTestId("model-pause-toggle-model-1"));
@ -368,7 +503,7 @@ describe("AllModelsTab", () => {
it("opens the model settings modal from the toolbar", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(screen.queryByTestId("model-settings-modal")).not.toBeInTheDocument();
await user.click(screen.getByTestId("models-settings-trigger"));
@ -377,7 +512,7 @@ describe("AllModelsTab", () => {
it("opens the model detail view from the model ID cell", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(await screen.findByTestId("model-id-model-1"));
@ -386,7 +521,7 @@ describe("AllModelsTab", () => {
it("opens the team detail view from the team ID cell", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(await screen.findByTestId("model-team-id-model-1"));
@ -395,20 +530,20 @@ describe("AllModelsTab", () => {
describe("virtual key hint", () => {
it("explains personal key creation while viewing current team models", () => {
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(screen.getByText(/create a Virtual Key without selecting a team/i)).toBeInTheDocument();
});
it("links the Virtual Keys page through the migrated /ui route", () => {
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(screen.getByRole("link", { name: "Virtual Keys page" })).toHaveAttribute("href", "/ui/api-keys");
});
it("links the team hint's Virtual Keys page through the migrated /ui route", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(screen.getByTestId("models-team-select"));
await user.click(await screen.findByRole("option", { name: "Engineering" }));
@ -419,7 +554,7 @@ describe("AllModelsTab", () => {
it("names the selected team in the hint", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(screen.getByTestId("models-team-select"));
await user.click(await screen.findByRole("option", { name: "Engineering" }));
@ -429,7 +564,7 @@ describe("AllModelsTab", () => {
it("hides the hint when viewing all available models", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(screen.getByTestId("models-view-select"));
await user.click(await screen.findByRole("option", { name: "All Available Models" }));

View file

@ -10,10 +10,11 @@ import { toast } from "@/lib/toast";
import { uiHref } from "@/utils/uiHref";
import { modelDeleteCall, modelPatchUpdateCall } from "@/components/networking";
import { useQueryClient } from "@tanstack/react-query";
import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer";
import { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table";
import { useDebouncedValue } from "@tanstack/react-pacer/debouncer";
import { ColumnFiltersState, functionalUpdate, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table";
import { Info } from "lucide-react";
import { useCallback, useEffect, useMemo, useState } from "react";
import { createParser, parseAsInteger, parseAsString, parseAsStringLiteral, useQueryStates } from "nuqs";
import { useCallback, useMemo, useState } from "react";
import { useModelsInfo } from "../../hooks/models/useModels";
import { transformModelData } from "../utils/modelDataTransformer";
@ -24,11 +25,40 @@ import {
PERSONAL_TEAM_VALUE,
WILDCARD_MODEL_GROUP_VALUE,
} from "./AllModelsTable";
import { ACCESS_GROUPS_COLUMN_ID, MODEL_NAME_COLUMN_ID, toServerSortField } from "./ModelsTableColumns";
import {
ACCESS_GROUPS_COLUMN_ID,
isModelTableSortColumnId,
MODEL_NAME_COLUMN_ID,
MODEL_TABLE_SORT_COLUMN_IDS,
toServerSortField,
} from "./ModelsTableColumns";
const SEARCH_DEBOUNCE_WAIT_MS = 200;
const DEFAULT_PAGE_SIZE = 50;
const DEFAULT_PAGINATION: PaginationState = { pageIndex: 0, pageSize: DEFAULT_PAGE_SIZE };
const MAX_PAGE_SIZE = 100;
const MAX_PAGE = 100_000;
const MODEL_VIEW_MODES = ["current_team", "all"] as const satisfies readonly ModelViewMode[];
const boundedInteger = (min: number, max: number, fallback: number) =>
createParser({
parse: (value: string) => {
const parsed = parseAsInteger.parse(value);
return parsed === null ? null : Math.min(Math.max(parsed, min), max);
},
serialize: String,
}).withDefault(fallback);
const TABLE_STATE = {
model_search: parseAsString.withDefault(""),
view_mode: parseAsStringLiteral(MODEL_VIEW_MODES).withDefault("current_team"),
filter_team: parseAsString.withDefault(PERSONAL_TEAM_VALUE),
access_group: parseAsString.withDefault(""),
sort_by: parseAsStringLiteral(MODEL_TABLE_SORT_COLUMN_IDS),
sort_order: parseAsStringLiteral(["asc", "desc"] as const).withDefault("asc"),
page: boundedInteger(1, MAX_PAGE, 1),
page_size: boundedInteger(1, MAX_PAGE_SIZE, DEFAULT_PAGE_SIZE),
};
interface AllModelsTabProps {
selectedModelGroup: string | null;
@ -52,34 +82,25 @@ const AllModelsTab = ({
const { data: teams, isLoading: isLoadingTeams } = useTeams();
const queryClient = useQueryClient();
const [modelNameSearch, setModelNameSearch] = useState<string>("");
const [debouncedSearch, setDebouncedSearch] = useState<string>("");
const [modelViewMode, setModelViewMode] = useState<ModelViewMode>("current_team");
const [selectedTeamValue, setSelectedTeamValue] = useState<string>(PERSONAL_TEAM_VALUE);
const [selectedModelAccessGroupFilter, setSelectedModelAccessGroupFilter] = useState<string | null>(null);
const [pagination, setPagination] = useState<PaginationState>(DEFAULT_PAGINATION);
const [sorting, setSorting] = useState<SortingState>([]);
const [tableState, setTableState] = useQueryStates(TABLE_STATE);
const modelNameSearch = tableState.model_search;
const [debouncedSearch] = useDebouncedValue(modelNameSearch, { wait: SEARCH_DEBOUNCE_WAIT_MS });
const modelViewMode = tableState.view_mode;
const selectedTeamValue = tableState.filter_team;
const selectedModelAccessGroupFilter = tableState.access_group || null;
const pagination = useMemo<PaginationState>(
() => ({ pageIndex: tableState.page - 1, pageSize: tableState.page_size }),
[tableState.page, tableState.page_size],
);
const sorting = useMemo<SortingState>(
() => (tableState.sort_by ? [{ id: tableState.sort_by, desc: tableState.sort_order === "desc" }] : []),
[tableState.sort_by, tableState.sort_order],
);
const [isModelSettingsModalVisible, setIsModelSettingsModalVisible] = useState(false);
const [deleteModalModelId, setDeleteModalModelId] = useState<string | null>(null);
const [deleteLoading, setDeleteLoading] = useState(false);
const [pausingModelId, setPausingModelId] = useState<string | null>(null);
const resetToFirstPage = useCallback(() => {
setPagination((previous) => (previous.pageIndex === 0 ? previous : { ...previous, pageIndex: 0 }));
}, []);
const debouncedUpdateSearch = useDebouncedCallback(
(value: string) => {
setDebouncedSearch(value);
resetToFirstPage();
},
{ wait: SEARCH_DEBOUNCE_WAIT_MS },
);
useEffect(() => {
debouncedUpdateSearch(modelNameSearch);
}, [modelNameSearch, debouncedUpdateSearch]);
const teamIdForQuery = selectedTeamValue === PERSONAL_TEAM_VALUE ? undefined : selectedTeamValue;
const isConcreteModelGroup =
Boolean(selectedModelGroup) &&
@ -152,33 +173,49 @@ const AllModelsTab = ({
[selectedModelGroup, selectedModelAccessGroupFilter],
);
const handleSearchChange = useCallback(
(value: string) => {
void setTableState({ model_search: value || null, page: null });
},
[setTableState],
);
const handleColumnFiltersChange: OnChangeFn<ColumnFiltersState> = (updater) => {
const next = typeof updater === "function" ? updater(columnFilters) : updater;
const next = functionalUpdate(updater, columnFilters);
const modelGroup = next.find((entry) => entry.id === MODEL_NAME_COLUMN_ID)?.value;
const accessGroup = next.find((entry) => entry.id === ACCESS_GROUPS_COLUMN_ID)?.value;
setSelectedModelGroup(typeof modelGroup === "string" ? modelGroup : ALL_MODEL_GROUPS_VALUE);
setSelectedModelAccessGroupFilter(typeof accessGroup === "string" ? accessGroup : null);
resetToFirstPage();
void setTableState({ access_group: typeof accessGroup === "string" ? accessGroup : null, page: null });
};
const handleSortingChange: OnChangeFn<SortingState> = (updater) => {
setSorting(typeof updater === "function" ? updater(sorting) : updater);
resetToFirstPage();
const active = functionalUpdate(updater, sorting)[0];
void setTableState({
sort_by: active && isModelTableSortColumnId(active.id) ? active.id : null,
sort_order: active?.desc ? "desc" : null,
page: null,
});
};
const handlePaginationChange = useCallback<OnChangeFn<PaginationState>>(
(updater) => {
const next = functionalUpdate(updater, pagination);
void setTableState({ page: next.pageIndex + 1, page_size: next.pageSize });
},
[pagination, setTableState],
);
const handleTeamChange = (value: string) => {
setSelectedTeamValue(value);
resetToFirstPage();
void setTableState({ filter_team: value, page: null });
};
const handleViewModeChange = (value: ModelViewMode) => {
void setTableState({ view_mode: value });
};
const resetFilters = () => {
setModelNameSearch("");
setSelectedModelGroup(ALL_MODEL_GROUPS_VALUE);
setSelectedModelAccessGroupFilter(null);
setSelectedTeamValue(PERSONAL_TEAM_VALUE);
setModelViewMode("current_team");
setPagination(DEFAULT_PAGINATION);
setSorting([]);
void setTableState(null);
};
const teamOptions = useMemo(
@ -264,18 +301,18 @@ const AllModelsTab = ({
sorting={sorting}
onSortingChange={handleSortingChange}
pagination={pagination}
onPaginationChange={setPagination}
onPaginationChange={handlePaginationChange}
columnFilters={columnFilters}
onColumnFiltersChange={handleColumnFiltersChange}
onResetFilters={resetFilters}
searchValue={modelNameSearch}
onSearchChange={setModelNameSearch}
onSearchChange={handleSearchChange}
teamOptions={teamOptions}
selectedTeamValue={selectedTeamValue}
onTeamChange={handleTeamChange}
isLoadingTeams={isLoadingTeams}
viewMode={modelViewMode}
onViewModeChange={setModelViewMode}
onViewModeChange={handleViewModeChange}
onOpenModelSettings={handleOpenModelSettings}
availableModelGroups={availableModelGroups}
availableModelAccessGroups={availableModelAccessGroups}

View file

@ -24,6 +24,19 @@ export const TEAM_ID_COLUMN_ID = "model_info_team_id";
export const ACCESS_GROUPS_COLUMN_ID = "model_info_access_groups";
export const STATUS_COLUMN_ID = "model_info_db_model";
export const MODEL_TABLE_SORT_COLUMN_IDS = [
MODEL_NAME_COLUMN_ID,
CREATED_BY_COLUMN_ID,
UPDATED_AT_COLUMN_ID,
COSTS_COLUMN_ID,
STATUS_COLUMN_ID,
] as const;
export type ModelTableSortColumnId = (typeof MODEL_TABLE_SORT_COLUMN_IDS)[number];
export const isModelTableSortColumnId = (columnId: string): columnId is ModelTableSortColumnId =>
(MODEL_TABLE_SORT_COLUMN_IDS as readonly string[]).includes(columnId);
const COLUMN_ID_TO_SERVER_SORT_FIELD: Record<string, string> = {
[COSTS_COLUMN_ID]: "costs",
[STATUS_COLUMN_ID]: "status",

View file

@ -35,10 +35,12 @@ beforeEach(() => {
Element.prototype.scrollIntoView = () => {};
});
const CHAT_REQUEST_ARG_COUNT = 26;
const CHAT_REQUEST_ARG_COUNT = 27;
const STREAMING_ENABLED_ARG_INDEX = 25;
const MESSAGES_REQUEST_ARG_COUNT = 19;
const CHAT_CUSTOM_HEADERS_ARG_INDEX = 26;
const MESSAGES_REQUEST_ARG_COUNT = 20;
const MESSAGES_STREAMING_ENABLED_ARG_INDEX = 18;
const MESSAGES_CUSTOM_HEADERS_ARG_INDEX = 19;
async function openComboboxByPlaceholder(placeholder: string) {
const user = userEvent.setup();
@ -447,6 +449,63 @@ describe("ChatUI", () => {
expect(requestArgs[MESSAGES_STREAMING_ENABLED_ARG_INDEX]).toBe(false);
});
it("should send custom headers entered in the sidebar with /v1/chat/completions and /v1/messages requests", async () => {
const user = userEvent.setup();
render(
<ChatUI
accessToken="1234567890"
token="1234567890"
userRole="user"
userID="1234567890"
disabledPersonalKeyCreation={false}
/>,
);
await waitFor(() => {
expect(screen.getByText("Test Key")).toBeInTheDocument();
});
await selectComboboxOption("Select a Model", "Model 1");
await user.click(screen.getByRole("button", { name: "Add Header" }));
await user.click(screen.getByRole("button", { name: "Add Header" }));
const [firstName] = screen.getAllByPlaceholderText("Header Name");
const [firstValue, secondValue] = screen.getAllByPlaceholderText("Header Value");
fireEvent.change(firstName, { target: { value: "anthropic-beta" } });
fireEvent.change(firstValue, { target: { value: "context-1m-2025-08-07" } });
fireEvent.change(secondValue, { target: { value: "ignored because the name is blank" } });
const messageInput = screen.getByPlaceholderText("Type your message... (Shift+Enter for new line)");
await act(async () => {
fireEvent.change(messageInput, { target: { value: "hello" } });
});
await act(async () => {
fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" });
});
await waitFor(() => {
expect(makeOpenAIChatCompletionRequest).toHaveBeenCalledTimes(1);
});
const chatArgs = vi.mocked(makeOpenAIChatCompletionRequest).mock.calls[0];
expect(chatArgs).toHaveLength(CHAT_REQUEST_ARG_COUNT);
expect(chatArgs[CHAT_CUSTOM_HEADERS_ARG_INDEX]).toEqual({ "anthropic-beta": "context-1m-2025-08-07" });
await selectComboboxOption("Select an endpoint", "/v1/messages");
await selectComboboxOption("Select a Model", "Model 1");
await act(async () => {
fireEvent.change(messageInput, { target: { value: "hello again" } });
});
await act(async () => {
fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" });
});
await waitFor(() => {
expect(makeAnthropicMessagesRequest).toHaveBeenCalledTimes(1);
});
const messagesArgs = vi.mocked(makeAnthropicMessagesRequest).mock.calls[0];
expect(messagesArgs).toHaveLength(MESSAGES_REQUEST_ARG_COUNT);
expect(messagesArgs[MESSAGES_CUSTOM_HEADERS_ARG_INDEX]).toEqual({ "anthropic-beta": "context-1m-2025-08-07" });
});
it("should force streaming in simplified mode even when the playground setting is off", async () => {
sessionStorage.setItem("streamingEnabled", "false");

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