diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index 7a7d1d776d3..46c90295198 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -125,6 +125,8 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): status_code=e.response.status_code, message=e.response.text, ) + except litellm.Timeout as e: + raise AnthropicError(status_code=e.status_code, message=str(e)) from e except Exception as e: verbose_logger.error("Error in CountTokens handler: %s", e) raise AnthropicError( diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 99b1cd7f241..755ccd129d8 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -18,7 +18,7 @@ from litellm.llms.anthropic.count_tokens.transformation import ( ChatCountTranslationError, anthropic_count_fields, ) -from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.base_llm.base_utils import BaseTokenCounter, deployment_timeout from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests @@ -134,6 +134,7 @@ class AnthropicTokenCounter(BaseTokenCounter): optional_params=count_options, context_management=context_management.spec if context_management is not None else None, drop_params=drops_unsupported_params(litellm_params.get("drop_params")), + timeout=deployment_timeout(litellm_params), ) if result is not None: diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 5fcb84593d2..443a3a5edad 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -147,6 +147,8 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): status_code=e.response.status_code, message=e.response.text, ) + except litellm.Timeout as e: + raise AnthropicError(status_code=e.status_code, message=str(e)) from e except Exception as e: verbose_logger.error("Error in CountTokens handler: %s", e) raise AnthropicError( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index 8657fcddab4..5031b69f83b 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -20,7 +20,7 @@ from litellm.llms.anthropic.count_tokens.transformation import ( from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) -from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.base_llm.base_utils import BaseTokenCounter, deployment_timeout from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests @@ -137,6 +137,7 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): system=system, optional_params=count_options, context_management=None if context_management is None else context_management.spec, + timeout=deployment_timeout(litellm_params), ) if result is not None: diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index 45f623730a8..8b1ccda3a61 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -15,8 +15,9 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Any, Final, TypeAlias +import httpx from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # SDK parser internals -from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, JsonValue, PositiveFloat, TypeAdapter, ValidationError from typing_extensions import TypeIs # noqa: TID251 # narrows a message to its system member by the role literal from litellm._logging import verbose_logger @@ -38,6 +39,18 @@ from litellm.types.utils import ( TokenCountResponse, ) +_SECONDS: Final[TypeAdapter[float]] = TypeAdapter(PositiveFloat) + + +def deployment_timeout(litellm_params: Mapping[str, object]) -> float | httpx.Timeout | None: + timeout: Final = litellm_params.get("timeout") + if timeout is None or isinstance(timeout, httpx.Timeout): + return timeout + try: + return _SECONDS.validate_python(timeout) + except ValidationError: + return None + class BaseTokenCounter(ABC): @abstractmethod diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 2cf48a75524..ce9ed6c0940 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -111,6 +111,8 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): status_code=e.response.status_code, message=e.response.text, ) + except litellm.Timeout as e: + raise OpenAIError(status_code=e.status_code, message=str(e)) from e except (httpx.RequestError, json.JSONDecodeError, ValueError) as e: verbose_logger.error("Error in CountTokens handler: %s", e) raise OpenAIError( diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index c82a2982a54..cd60bb1327f 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Any, Final +import httpx import openai from pydantic import JsonValue from typing_extensions import assert_never @@ -21,7 +22,7 @@ from litellm.llms.anthropic.messages.responses_adapters.handler import count_vie from litellm.llms.anthropic.messages.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, ) -from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.base_llm.base_utils import BaseTokenCounter, deployment_timeout from litellm.llms.openai.common_utils import OpenAIError, decoded_json_object, get_openai_credentials from litellm.llms.openai.responses.count_tokens.handler import ( OpenAICountTokensHandler, @@ -61,6 +62,7 @@ def _workload_identity_config( class _CountCredential: api_base: str | None credential: str | OpenAIWorkloadIdentityConfig + timeout: float | httpx.Timeout | None def _count_credential(deployment: Mapping[str, object] | None, model_label: str) -> _CountCredential | OpenAIError: @@ -82,7 +84,9 @@ def _count_credential(deployment: Mapping[str, object] | None, model_label: str) "workload identity on the deployment, or OPENAI_API_KEY in the proxy environment." ), ) - return _CountCredential(api_base=credentials.api_base, credential=credential) + return _CountCredential( + api_base=credentials.api_base, credential=credential, timeout=deployment_timeout(litellm_params) + ) def _count_error(model_to_use: str, request_model: str, message: str, status_code: int) -> TokenCountResponse: @@ -245,6 +249,7 @@ class OpenAITokenCounter(BaseTokenCounter): tools=counted_tools, instructions=instructions, text=text, + timeout=resolved.timeout, ), model_to_use, request_model, @@ -264,7 +269,9 @@ class OpenAITokenCounter(BaseTokenCounter): body: Final = native_count_request(request, model_to_use, policy) return await self._counted( resolved, - lambda api_key, api_base: openai_count_tokens_handler.post_count_request(body, api_key, api_base), + lambda api_key, api_base: openai_count_tokens_handler.post_count_request( + body, api_key, api_base, resolved.timeout + ), model_to_use, request_model, ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3196b3c0c08..a4c16fbeaa3 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3965,6 +3965,7 @@ "supports_vision": true, "supports_output_config": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 4096 }, "azure_ai/claude-opus-4-7": { @@ -3997,6 +3998,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 2048 }, "azure_ai/claude-fable-5": { @@ -4031,6 +4033,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": false, + "supports_web_search": true, "prompt_cache_min_tokens": 512 }, "azure_ai/claude-fable-5-1": { @@ -4065,6 +4068,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": false, + "supports_web_search": true, "prompt_cache_min_tokens": 512, "deprecation_date": "2027-12-05", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement" @@ -4136,6 +4140,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 512 }, "azure_ai/claude-opus-5-5": { @@ -4171,6 +4176,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 512, "source": "https://management.azure.com/subscriptions/c873328e-b572-4770-8dff-aaeb6f1f0e79/providers/Microsoft.CognitiveServices/locations/eastus2/models?api-version=2024-10-01" }, @@ -4208,6 +4214,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 1024 }, "azure_ai/claude-sonnet-4-5": { @@ -4267,6 +4274,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 1024 }, "azure_ai/claude-sonnet-4-6": { @@ -72516,6 +72524,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.4-mini": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 8.25e-08, @@ -72537,6 +72546,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.4-nano": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 2.2e-08, @@ -73140,6 +73150,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.4-mini": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 8.25e-08, @@ -73161,6 +73172,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.4-nano": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 2.2e-08, @@ -76633,6 +76645,7 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, "supports_forced_tool_use": false, + "supports_web_search": true, "thinking_always_on": true, "prompt_cache_min_tokens": 512, "source": "https://management.azure.com/subscriptions/c873328e-b572-4770-8dff-aaeb6f1f0e79/providers/Microsoft.CognitiveServices/locations/eastus2/models?api-version=2024-10-01" @@ -78204,6 +78217,7 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, "supports_forced_tool_use": true, + "supports_web_search": true, "thinking_always_on": false, "prompt_cache_min_tokens": 512, "source": "https://platform.claude.com/docs/en/models/haiku-5-5/migration-guide" @@ -78910,6 +78924,7 @@ "prompt_cache_min_tokens": 4096.0, "rpm": 2000.0, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -78970,6 +78985,7 @@ "prompt_cache_min_tokens": 6144.0, "regional_endpoint_uplift_multiplier": 1.1, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -79029,6 +79045,7 @@ "prompt_cache_min_tokens": 6144.0, "regional_endpoint_uplift_multiplier": 1.1, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -79270,6 +79287,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "baseten/meta-models/muse-spark-open": { @@ -88127,6 +88145,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "azure_ai/claude-mythos-5-2-fast": { @@ -88162,6 +88181,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "azure_ai/claude-fable-5-2": { @@ -88193,6 +88213,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "azure_ai/claude-fable-5-2-fast": { @@ -88224,6 +88245,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "bedrock/us-gov-east-1/anthropic.claude-mythos-5-2": { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index abf61a909b2..228e80da4bd 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1860,6 +1860,21 @@ class CustomerBase(LiteLLMPydanticObjectBase): blocked: bool = False +def reject_unknown_customer_fields(request: type[BaseModel], values: object) -> object: + if not isinstance(values, dict): + return values + unknown: Final = tuple(sorted(str(key) for key in values if key not in request.model_fields)) # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # raw request body + if not unknown: + return values + on_budget: Final = tuple(field for field in unknown if field in BudgetNewRequest.model_fields) + hint: Final = ( + f" {', '.join(on_budget)}: these live on the customer's budget object, update it with /budget/update." + if on_budget + else "" + ) + raise ValueError(f"Unknown fields: {', '.join(unknown)}.{hint}") + + class NewCustomerRequest(BudgetNewRequest): """ Create a new customer, allocate a budget to them @@ -1877,6 +1892,11 @@ class NewCustomerRequest(BudgetNewRequest): models: list[str] | None = None object_permission: LiteLLM_ObjectPermissionBase | None = None + @model_validator(mode="before") + @classmethod + def reject_unknown_fields(cls, values: object) -> object: + return reject_unknown_customer_fields(cls, values) + @model_validator(mode="before") @classmethod def check_user_info(cls, values): @@ -1904,6 +1924,11 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): models: list[str] | None = None object_permission: LiteLLM_ObjectPermissionBase | None = None + @model_validator(mode="before") + @classmethod + def reject_unknown_fields(cls, values: object) -> object: + return reject_unknown_customer_fields(cls, values) + class DeleteCustomerRequest(LiteLLMPydanticObjectBase): """ diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index aec5e6840a5..bd202b85f73 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -353,7 +353,6 @@ async def new_end_user( - allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] - Require all user requests to use models in this specific region. - default_model: Optional[str] - If no equivalent model in the allowed region, default all requests to this model. - models: Optional[list[str]] - Restrict this customer's access to the listed models. - - metadata: Optional[dict] = Metadata for customer, store information for customer. Example metadata = {"data_training_opt_out": True} - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - tpm_limit: Optional[int] - [Not Implemented Yet] Specify tpm limit for a given customer (Tokens per minute) - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute) @@ -387,7 +386,7 @@ async def new_end_user( --header 'Content-Type: application/json' \ --data '{ "user_id" : "ishaan-jaff-3", - "allowed_region": "eu", + "allowed_model_region": "eu", "budget_id": "free_tier", "models": ["gpt-4o-mini"], "default_model": "azure/gpt-3.5-turbo-eu" diff --git a/litellm/proxy/mcp_server/discoverable_endpoints.py b/litellm/proxy/mcp_server/discoverable_endpoints.py index 07a267d6caa..ecc40b6e98f 100644 --- a/litellm/proxy/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/mcp_server/discoverable_endpoints.py @@ -105,9 +105,11 @@ from litellm.proxy.mcp_server.oauth_utils import ( get_cimd_client_id, get_cimd_document_url, get_request_base_url, + grant_covers_scopes, needs_cimd_discovery, oauth_client_registration_matches, resolve_upstream_resource, + server_scope_names, validate_trusted_redirect_uri, well_known_root_suffix, ) @@ -2537,6 +2539,9 @@ async def _vendor_credential_state( get_user_oauth_credential, oauth_grant_state, ) + from litellm.proxy.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # circular import at module load + global_mcp_server_manager, + ) from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # circular import at module load if prisma_client is None: @@ -2548,7 +2553,12 @@ async def _vendor_credential_state( if oauth_grant_state(credential) == "absent": return "absent" granted_scopes: Final = credential.get("scopes") if credential else None - return "present" if set(requested_scopes).issubset(granted_scopes or ()) else "absent" + covered: Final = grant_covers_scopes( + tuple(granted_scopes or ()), + requested_scopes, + lambda: server_scope_names(global_mcp_server_manager.get_mcp_server_by_id(server_id)), + ) + return "present" if covered else "absent" @router.get("/authorize/flow") diff --git a/litellm/proxy/mcp_server/oauth_utils.py b/litellm/proxy/mcp_server/oauth_utils.py index 76343f77e32..3171c476ede 100644 --- a/litellm/proxy/mcp_server/oauth_utils.py +++ b/litellm/proxy/mcp_server/oauth_utils.py @@ -3,7 +3,7 @@ import os import re -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass from ipaddress import ip_address from types import MappingProxyType @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Final, NoReturn from urllib.parse import ParseResult, parse_qsl, urlencode, urlparse, urlsplit, urlunparse, urlunsplit from fastapi import HTTPException, Request +from pydantic import TypeAdapter from starlette.types import Scope from litellm._logging import verbose_logger @@ -131,6 +132,34 @@ def build_upstream_authorize_url(authorization_url: str, params: Mapping[str, st return urlunparse(parsed._replace(query=urlencode(merged))) +def grant_covers_scopes( + granted: tuple[str, ...], requested: tuple[str, ...], server_scopes: Callable[[], frozenset[str]] +) -> bool: + """Whether a stored grant holds every requested scope. A grant stored as one comma-joined token, the form + Slack and GitHub reply with, counts as its parts only when every part is a requested or server scope.""" + if all(scope in granted for scope in requested): + return True + if len(granted) != 1 or "," not in granted[0]: + return False + parts: Final = frozenset(granted[0].split(",")) + known: Final = frozenset(requested) | server_scopes() + return parts <= known and all(scope in granted or scope in parts for scope in requested) + + +_PUBLISHED_SCOPES: Final = TypeAdapter[tuple[str, ...] | None](tuple[str, ...] | None) + + +def server_scope_names(server: "MCPServer | None") -> frozenset[str]: + if server is None: + return frozenset() + published: Final = ( + _PUBLISHED_SCOPES.validate_python(server.authorization_server_metadata.get("scopes_supported")) + if server.authorization_server_metadata is not None + else None + ) + return frozenset((*(server.scopes or ()), *(published or ()))) + + def proxy_base_url_from_env() -> str | None: """The gateway's fixed public URL from ``PROXY_BASE_URL``, ``None`` when unset or malformed.""" return _resolve_proxy_base_url_env() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e6447754436..3fddca6f4d8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15325,7 +15325,12 @@ def _provider_count_outcome( if result is not None and result.error is True: return ProviderCountFailure( status_code=result.status_code or 500, - message=result.error_message or "Token counting failed", + message=result.error_message + or ( + "Token counting failed" + if result.status_code is None + else f"The provider's count API answered {result.status_code} with no error message" + ), rejected_before_provider_call=result.rejected_before_provider_call, ) return result @@ -15794,6 +15799,14 @@ def _token_count_settings_model(deployment: _TokenCountDeployment, request_model _PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset( {litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value} ) +_CHAT_MODES: Final = frozenset({"chat", "responses"}) + + +def _serves_chat(deployment: _TokenCountDeployment, cost_map_row: Mapping[str, object]) -> bool: + mode: Final = _non_empty_str(deployment.get("model_info", {}).get("mode")) or _non_empty_str( + cost_map_row.get("mode") + ) + return mode is None or mode in _CHAT_MODES _ESTIMABLE_GOOGLE_PART_KEYS: Final = frozenset({"text", "functionCall", "functionResponse", "thoughtSignature"}) @@ -16104,6 +16117,12 @@ async def _routed_count_deployment( return None +def _with_chat_timeout(litellm_params: Mapping[str, object]) -> Mapping[str, object]: + if llm_router is None: + return litellm_params + return {**litellm_params, "timeout": llm_router.settings().timeout_for(litellm_params)} + + async def _count_request_tokens( request: TokenCountRequest, routed: Mapping[str, object] | None, @@ -16128,8 +16147,8 @@ async def _count_request_tokens( litellm_model_name = deployment.get("litellm_params", {}).get("model") deployment = { # rebind-ok: merge the named credential into a copy **deployment, - "litellm_params": with_named_credential( - _LITELLM_PARAMS_COPY.validate_python(deployment.get("litellm_params", {})) + "litellm_params": _with_chat_timeout( + with_named_credential(_LITELLM_PARAMS_COPY.validate_python(deployment.get("litellm_params", {}))) ), } # remove the custom_llm_provider_prefix in the litellm_model_name @@ -16142,9 +16161,10 @@ async def _count_request_tokens( deployment_view: Final = _TOKEN_COUNT_DEPLOYMENT.validate_python( deployment or _model_group_deployment(llm_router, request.model) or {} ) + cost_map_row: Final = await asyncify(cost_map_entry)(*_token_count_settings_model(deployment_view, request.model)) settings: Final = await asyncify(resolve_token_count_settings)( configured=deployment_view.get("model_info", {}), - default=await asyncify(cost_map_entry)(*_token_count_settings_model(deployment_view, request.model)), + default=cost_map_row, deployment_region=partial(_deployment_aws_region, deployment_view), ) @@ -16169,6 +16189,15 @@ async def _count_request_tokens( code=400, ) prompt_only: Final = bool(prompt) and conversation_empty + prompt_as_turn: Final = ( + prompt_only + and provider_counter is not None + and provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider) + and _serves_chat(deployment_view, cost_map_row) + ) + prompt_as_message: Final = prompt_as_turn or ( + prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES + ) request_messages: Final = _counted_request_messages(messages, messages_format, count_deployment) local_count_error: Final = malformed_tool_result_count_error(request_messages, messages_format, count_deployment) malformed_messages: Final = ( @@ -16205,7 +16234,6 @@ async def _count_request_tokens( ) inference_system: Final = system if edited_request is None else edited_request.system contents_estimable: Final = not contents or (bool(messages) and _estimable_contents(contents)) - prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES routed_deployment: Final = _ROUTED_DEPLOYMENT.validate_python(deployment) routed_request: Final = ( None if responses_request is None else with_provider_file_ids(responses_request, routed_deployment) @@ -16306,7 +16334,7 @@ async def _count_request_tokens( _with_get(model_info).get("custom_tokenizer", None), ) system_message: Final = _system_message(local_system) - prompt_turn: Final = prompt_only and not (system_message is None and not tools) + prompt_turn: Final = prompt_as_turn or (prompt_only and not (system_message is None and not tools)) typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes Sequence[AllMessageValues] | None, [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else local_messages), diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index 683a51fc6b6..6a55302027d 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -192,35 +192,11 @@ def responses_api_bridge_check( model_info["mode"] = "responses" model = model.replace("responses/", "") - # OpenAI/Azure GPT-5 chat-completions that need Responses-only fields (e.g. - # ``reasoningSummary`` in ``extra_body``) must be bridged; Chat Completions rejects - # those keys. - # - # - gpt-5.4+: FUNCTION tools with reasoning active must be bridged. OpenAI enables - # reasoning by default for these models (unset reasoning_effort means medium - # server-side), and Chat Completions rejects function tools whenever reasoning is - # on ("Function tools with reasoning_effort are not supported ... use - # /v1/responses or set reasoning_effort to 'none'"), so only an explicit - # ``"none"`` the chat gate forwards keeps the request chat-servable: a ``none`` - # the model's map row refuses is dropped (or refused) before the call, which - # leaves the provider's default effort in force. Custom (grammar) tools are served - # natively by Chat Completions with reasoning on, so custom-only requests stay on - # chat and keep their native custom tool_call response shape. - # - The UNSET-effort arm only fires against endpoints known to enforce that - # constraint (any api.openai.com host, or Azure OpenAI where api_base is - # always set): chat-only OpenAI-compatible backends registered under the openai - # provider with a custom api_base and gpt-5.4+ model names serve tools without - # reasoning fine and have no /responses route, so they keep pre-existing - # behavior (bridge only on an explicit reasoning_effort). - # - Azure AI Foundry's OpenAI v1 hosts (azure_ai provider) enforce it later in the series: - # an explicit effort with function tools is rejected from gpt-5.6 on, and the unset - # effort only from gpt-6 on (gpt-5.6 serves tools with reasoning silently off), so the - # azure_ai gate keys on those measured boundaries instead of gpt-5.4+. - # - Bedrock Mantle serves the same API and enforces the gpt-5.4+ rule even with the effort - # unset (measured on gpt-5.6-terra), so a Mantle row that also serves Responses follows - # these rules, which keeps a deployment pinned to ``mode: chat`` from a 400 on them. - # - Older GPT-5 names (e.g. ``gpt-5``, ``gpt-5.1``): bridge only when a reasoning - # summary alias is present with ``reasoning_effort`` (tools alone stay on chat). + # Function tools with explicit active reasoning bridge on GPT-5.4+. With effort unset, + # OpenAI/Azure and Responses-capable Mantle endpoints bridge only GPT-5.6+ whose map + # default is not "none". Custom OpenAI-compatible endpoints keep unset effort on chat. + # Foundry uses its separate measured boundaries. Custom-only tools remain chat-native; + # reasoning summaries and tool_search retain their independent Responses requirements. has_function_tool: Final = any( ( tool.get("type") == "function" and (isinstance(tool.get("function"), dict) or "name" in tool) @@ -274,7 +250,14 @@ def responses_api_bridge_check( if on_foundry_openai_endpoint else ( OpenAIGPT5Config.is_model_gpt_5_4_plus_model(gpt_series_name) - and (reasoning_effort is not None or on_constraint_enforcing_endpoint) + and ( + reasoning_effort is not None + or ( + on_constraint_enforcing_endpoint + and OpenAIGPT5Config.is_model_gpt_5_6_plus_model(gpt_series_name) + and model_info.get("default_reasoning_effort") != "none" + ) + ) ) ) ) diff --git a/litellm/router_api.py b/litellm/router_api.py index d07b135e9a5..c85a764151f 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -1308,6 +1308,9 @@ class RouterSettings: """The guardrail configured under this name (was ``get_available_guardrail``).""" return self._owner("_core_get_available_guardrail")._core_get_available_guardrail(name) + def timeout_for(self, litellm_params: Mapping[str, object]) -> object: + return self._owner("_core_deployment_call_timeout")._core_deployment_call_timeout(litellm_params) + class RouterHealth: """Cooldowns, blocked deployments, health checks and usage (``health()``).""" diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 5015795657f..682aa557007 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -113,7 +113,7 @@ from litellm.router_rust.invoke.endpoints import ( stamp_error, unrouted_fields, ) -from litellm.router_rust.invoke.shaping import Shaping +from litellm.router_rust.invoke.shaping import Shaping, non_stream_timeout from litellm.router_rust.invoke.slots import enter_deployment_slot, max_parallel_requests_limit from litellm.router_rust.invoke.streaming import ( StreamAttempt, @@ -1592,6 +1592,9 @@ class RustCore( def _shaping(self) -> Shaping: return self._endpoints.shaping(self._resolved) + def _core_deployment_call_timeout(self, litellm_params: Mapping[str, object]) -> object: + return non_stream_timeout(self._shaping, MappingProxyType({}), litellm_params) + # -- building -------------------------------------------------------------------------------- def _new_handle(self, model_list: Sequence[object]) -> "RouterCoreHandle": @@ -1815,13 +1818,14 @@ class RustCore( def _rebuild(self) -> None: """A new Rust core for changed static settings, over the same deployments and the providers the old one's builds resolved for them.""" - old: Final = self._handle - model_list: Final = list(self._hosted._core_model_list) # pyright: ignore[reportPrivateUsage] # see __init__ - self._handle = self._new_handle(model_list) - self._handle.adopt_providers(old) - old.retire() - # The live list and indexes read the new handle's (the same entries, kept as the same objects). - self._hosted.refresh() + with self._hosted._deployment_index_lock: # pyright: ignore[reportPrivateUsage] # the registry's lock, held across the swap and the read back + old: Final = self._handle + model_list: Final = list(self._hosted._core_model_list) # pyright: ignore[reportPrivateUsage] # see __init__ + self._handle = self._new_handle(model_list) + self._handle.adopt_providers(old) + old.retire() + # The live list and indexes read the new handle's (the same entries, kept as the same objects). + self._hosted.refresh() # The new handle builds its indexes from the deployments: written ones are handed again. self._index_sent = None diff --git a/litellm/router_rust/surface/model_list.py b/litellm/router_rust/surface/model_list.py index bb7fb53acea..32ab8f1cb76 100644 --- a/litellm/router_rust/surface/model_list.py +++ b/litellm/router_rust/surface/model_list.py @@ -26,7 +26,9 @@ model info caches) mean what they did: the same object until a write changes tha import copy import datetime +import threading from collections.abc import Callable, Iterable, Mapping, Sequence +from contextlib import AbstractContextManager from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -189,9 +191,12 @@ class HostedModelList(list[Entry]): apply: Callable[[Mapping[str, object]], None], live: Callable[[object], Mapping[str, object]] = lambda _: {}, live_info: Callable[[object], Mapping[str, object]] = lambda _: {}, + *, + lock: AbstractContextManager[object] | None = None, ) -> None: super().__init__() self._handle: Final = handle + self._lock: Final[AbstractContextManager[object]] = threading.RLock() if lock is None else lock self._apply: Final = apply # A deployment's litellm_params values its record carries as text, by id (``live_params``). self._live: Final = live @@ -209,6 +214,10 @@ class HostedModelList(list[Entry]): def refresh(self) -> None: """Read router-core's deployments again when its registry moved since the last read.""" + with self._lock: + self._refresh() + + def _refresh(self) -> None: handle: Final = self._handle() if handle is None or self._building is not None: # While a build runs, the list shows what it accepted until router-core serves it. @@ -260,28 +269,31 @@ class HostedModelList(list[Entry]): """A build starts: router-core serves what it adds once it ends, so until then the list shows each deployment it accepts here (``built``), after the entries it keeps, or alone when it replaces them, as the Python router's list grows while it adds them.""" - self.refresh() - if replacing: - self._show([]) - self._id_counts = {} - self._building = (len(self), self._id_counts, self._read_at) + with self._lock: + self._refresh() + if replacing: + self._show([]) + self._id_counts = {} + self._building = (len(self), self._id_counts, self._read_at) def built(self, entry: Mapping[str, object]) -> None: """A deployment the build accepted, shown until the build ends.""" - if self._building is not None: - super().extend(self._kept([entry], appended=True)) + with self._lock: + if self._building is not None: + super().extend(self._kept([entry], appended=True)) def end_build(self) -> None: """The build ended (or failed): router-core serves what it added, read back in place of the entries shown meanwhile, each kept as the object it was shown as while unchanged.""" - if self._building is None: - return - start, counts, read_at = self._building - self._building = None - super().__delitem__(slice(start, None)) - self._id_counts = counts - self._read_at = read_at - self.refresh() + with self._lock: + if self._building is None: + return + start, counts, read_at = self._building + self._building = None + super().__delitem__(slice(start, None)) + self._id_counts = counts + self._read_at = read_at + self._refresh() def _seed(self, entries: Sequence[object]) -> None: """Keep ``entries``, the list router-core now holds, as the objects its records for them @@ -311,19 +323,21 @@ class HostedModelList(list[Entry]): def _replace(self, entries: list[Entry]) -> None: # mutable-ok: the list router-core takes # Seeded first: router-core's records for them read back as these objects however soon # the list is read again. - self._seed(entries) - self._apply({"replace": entries}) - self.refresh() + with self._lock: + self._seed(entries) + self._apply({"replace": entries}) + self._refresh() def replace(self, entries: Iterable[object]) -> None: """Every deployment at once (``router.model_list = [...]``).""" self._replace(cast(list[Entry], list(entries))) # cast-ok: the entries as assigned, which router-core validates def append(self, value: Entry) -> None: - self.refresh() - self._seed_appended([value]) - self._apply({"append": [value]}) - self.refresh() + with self._lock: + self._refresh() + self._seed_appended([value]) + self._apply({"append": [value]}) + self._refresh() def extend(self, values: Iterable[Entry]) -> None: for value in tuple(values): @@ -362,8 +376,9 @@ class HostedModelList(list[Entry]): # A delete removes every deployment with the id; the list removes only this one. self._replace(changed) return - self._apply({"delete": [removed_id]}) - self.refresh() + with self._lock: + self._apply({"delete": [removed_id]}) + self._refresh() def pop(self, index: SupportsIndex = -1) -> Entry: entry: Final = self[index] diff --git a/litellm/router_rust/surface/registry.py b/litellm/router_rust/surface/registry.py index e48569c4c9d..83008bdbb94 100644 --- a/litellm/router_rust/surface/registry.py +++ b/litellm/router_rust/surface/registry.py @@ -743,7 +743,7 @@ class HostRegistry(HostModelInfo): self._core_model_group_alias: dict[str, str | RouterModelGroupAliasItem] = model_group_alias # mutable-ok: the live alias map # fmt: skip # The live model_list: router-core's deployments read back, every write a changeset. self._core_model_list: Final[HostedModelList] = HostedModelList( - self._live_handle, self._hand_changes, self.live_params, self.live_info + self._live_handle, self._hand_changes, self.live_params, self.live_info, lock=self._deployment_index_lock ) # The litellm_params values a deployment's JSON cannot carry (an Exception or model object as # ``mock_response``, a callable ``azure_ad_token_provider``), by deployment id: an attempt is called with them, as Python's list held them. @@ -1472,9 +1472,10 @@ class HostRegistry(HostModelInfo): def _hand_changes(self, changes: Mapping[str, object]) -> None: if self.changes_runner is None: raise RuntimeError("the host registry writes deployments through its Rust core") - self.changes_runner(changes) - self._core_model_list.refresh() - self._drop_removed_live_values() + with self._deployment_index_lock: + self.changes_runner(changes) + self._core_model_list.refresh() + self._drop_removed_live_values() def _drop_removed_live_values(self) -> None: """Forget the live values of ids no deployment carries any more, so a router that rotates @@ -1678,10 +1679,11 @@ class HostRegistry(HostModelInfo): return handle is not None and key is not None and handle.has_model_id(key) def _core_get_model_info(self, id: str) -> UntypedEntry | None: - handle: Final = self._live_handle() - key: Final = _index_key(id) - idx: Final[object] = None if handle is None or key is None else handle.id_position(key) - return None if idx is None else self._core_model_list[idx] # pyright: ignore[reportArgumentType, reportCallIssue] # a position, or what a caller wrote, read as Python reads it + with self._deployment_index_lock: + handle: Final = self._live_handle() + key: Final = _index_key(id) + idx: Final[object] = None if handle is None or key is None else handle.id_position(key) + return None if idx is None else self._core_model_list[idx] # pyright: ignore[reportArgumentType, reportCallIssue] # a position, or what a caller wrote, read as Python reads it def live_params(self, model_id: object) -> Mapping[str, object]: """The litellm_params values of deployment ``model_id`` that its JSON record cannot carry.""" @@ -1696,21 +1698,23 @@ class HostRegistry(HostModelInfo): return None if entry is None else Deployment(**entry) # pyright: ignore[reportArgumentType] # a stored entry def _core_get_deployment_by_model_group_name(self, model_group_name: str) -> Deployment | None: - handle: Final = self._live_handle() - key: Final = _index_key(model_group_name) - indices: Final[object] = None if handle is None or key is None else handle.named_positions(key) - model: Final = self._core_model_list[indices[0]] if indices else None # pyright: ignore[reportIndexIssue, reportArgumentType, reportCallIssue] # positions, or what a caller wrote, read as Python reads them - return None if model is None else Deployment(**model) # pyright: ignore[reportArgumentType] # a stored entry + with self._deployment_index_lock: + handle: Final = self._live_handle() + key: Final = _index_key(model_group_name) + indices: Final[object] = None if handle is None or key is None else handle.named_positions(key) + model: Final = self._core_model_list[indices[0]] if indices else None # pyright: ignore[reportIndexIssue, reportArgumentType, reportCallIssue] # positions, or what a caller wrote, read as Python reads them + return None if model is None else Deployment(**model) # pyright: ignore[reportArgumentType] # a stored entry def _at(self, indices: object) -> tuple[UntypedEntry, ...]: """The entries at ``indices`` (what an index holds), read as Python's list reads them.""" return tuple(self._core_model_list[idx] for idx in indices) # pyright: ignore[reportGeneralTypeIssues, reportArgumentType, reportCallIssue] # positions, or what a caller wrote def _core_deployments_named(self, model_name: str) -> tuple[UntypedEntry, ...] | None: - handle: Final = self._live_handle() - key: Final = _index_key(model_name) - indices: Final[object] = None if handle is None or key is None else handle.named_positions(key) - return None if indices is None else self._at(indices) + with self._deployment_index_lock: + handle: Final = self._live_handle() + key: Final = _index_key(model_name) + indices: Final[object] = None if handle is None or key is None else handle.named_positions(key) + return None if indices is None else self._at(indices) def team_public_model_names_of(self, team_id: object) -> frozenset[str]: handle: Final = self._live_handle() @@ -1719,13 +1723,14 @@ class HostRegistry(HostModelInfo): return frozenset(handle.team_public_names_of(team_id)) def _team_deployments_named(self, team_id: str, public_model_name: str) -> tuple[UntypedEntry, ...] | None: - handle: Final = self._live_handle() - team: Final = _index_key(team_id) - public: Final = _index_key(public_model_name) - indices: Final[object] = ( - None if handle is None or team is None or public is None else handle.team_positions(team, public) - ) - return None if indices is None else self._at(indices) + with self._deployment_index_lock: + handle: Final = self._live_handle() + team: Final = _index_key(team_id) + public: Final = _index_key(public_model_name) + indices: Final[object] = ( + None if handle is None or team is None or public is None else handle.team_positions(team, public) + ) + return None if indices is None else self._at(indices) # -- lookups (router-core `lookup.rs`): rows read at the stored entries -------------------- @@ -1841,24 +1846,26 @@ class HostRegistry(HostModelInfo): def _team_deployments_across_teams(self, model: str) -> list[UntypedEntry]: # mutable-ok: a fresh list """Every team's deployments under public name ``model``, for a proxy admin calling without a team.""" - rows, teams = self._lookup().lookup_across_teams(model) - team_deployments: Final = self._rows(rows) - if teams > 1: - raise litellm_exceptions.BadRequestError( - message=( - f"Model name '{model}' matches deployments from multiple teams. " - "Specify the deployment ID directly to disambiguate." - ), - model=model, - llm_provider="", - ) - return team_deployments + with self._deployment_index_lock: + rows, teams = self._lookup().lookup_across_teams(model) + team_deployments: Final = self._rows(rows) + if teams > 1: + raise litellm_exceptions.BadRequestError( + message=( + f"Model name '{model}' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ), + model=model, + llm_provider="", + ) + return team_deployments def _core_get_all_deployments( self, model_name: str, model_alias: str | None = None, team_id: str | None = None ) -> list[UntypedEntry]: # mutable-ok: a fresh list, as the core returns it - key: Final = _index_key(model_name) - return [] if key is None else self._rows(self._lookup().lookup_all(key, team_id, model_alias)) + with self._deployment_index_lock: + key: Final = _index_key(model_name) + return [] if key is None else self._rows(self._lookup().lookup_all(key, team_id, model_alias)) def get_routing_group(self, model_name: str) -> RoutingGroup | None: """The routing group callable as ``model_name``: a deployment model_name or an alias of the @@ -1877,61 +1884,70 @@ class HostRegistry(HostModelInfo): def _core_routing_group_has_alternatives(self, model_group: str | None) -> bool: """A callable routing group whose member union spans more than one deployment.""" - if model_group is None: - return False - resolved: Final = self._core_get_model_from_alias(model=model_group) or model_group - group: Final = self.get_routing_group(resolved) - if group is None: - return False - handle: Final = self._lookup() + with self._deployment_index_lock: + if model_group is None: + return False + resolved: Final = self._core_get_model_from_alias(model=model_group) or model_group + group: Final = self.get_routing_group(resolved) + if group is None: + return False + handle: Final = self._lookup() - def held(member: str) -> Sequence[object]: - return cast(Sequence[object], handle.named_positions(member) or ()) # cast-ok: positions, or junk + def held(member: str) -> Sequence[object]: + return cast(Sequence[object], handle.named_positions(member) or ()) # cast-ok: positions, or junk - return sum(len(held(member)) for member in group.models) > 1 + return sum(len(held(member)) for member in group.models) > 1 def _core_get_model_list_from_model_alias( self, model_name: str | None = None ) -> list[UntypedEntry]: # mutable-ok: a fresh list - return self._rows( - self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="aliases")) - ) + with self._deployment_index_lock: + return self._rows( + self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="aliases")) + ) def _core_get_model_list_from_routing_groups(self, model_name: str | None = None) -> Sequence[UntypedEntry]: """Callable routing groups as model-list rows: each member deployment under the group's name.""" - if model_name is None and self._routing_group_rows is not None: - return self._routing_group_rows - rows: Final = tuple( - self._rows(self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="groups"))) - ) - if model_name is None: - self._routing_group_rows = rows - return rows + with self._deployment_index_lock: + if model_name is None and self._routing_group_rows is not None: + return self._routing_group_rows + rows: Final = tuple( + self._rows( + self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="groups")) + ) + ) + if model_name is None: + self._routing_group_rows = rows + return rows def _core_get_model_list( self, model_name: str | None = None, team_id: str | None = None ) -> list[UntypedEntry] | None: # mutable-ok: a fresh list, as get_model_list returns - if model_name is not None and _index_key(model_name) is None: - # No index holds it, so `_get_wildcard_deployments` reads it: what the pattern routers raise. - team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None - return [ - *self._core_pattern_router.get_deployments_by_pattern(model=model_name), - *(team_router.get_deployments_by_pattern(model=model_name) if team_router is not None else ()), - ] - return self._rows(self._lookup().lookup_model_list(model_name, team_id, self._list_scope(model_name))) + with self._deployment_index_lock: + if model_name is not None and _index_key(model_name) is None: + # No index holds it, so `_get_wildcard_deployments` reads it: what the pattern routers raise. + team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None + return [ + *self._core_pattern_router.get_deployments_by_pattern(model=model_name), + *(team_router.get_deployments_by_pattern(model=model_name) if team_router is not None else ()), + ] + return self._rows(self._lookup().lookup_model_list(model_name, team_id, self._list_scope(model_name))) def _core_wildcard_route_deployments(self, model: str, team_id: str | None) -> Sequence[Mapping[str, object]]: - from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + with self._deployment_index_lock: + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider - if not isinstance(model, str): # pyright: ignore[reportUnnecessaryIsInstance] # a caller's non-string model reads the base routers - team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None - team_owned: Final = team_router.get_deployments_by_pattern(model=model) if team_router is not None else () - return team_owned or self._core_pattern_router.get_deployments_by_pattern(model=model) - rows: Final = self._lookup().lookup_wildcard_route( - model, team_id, [model, f"{inferred_provider(model)}/{model}"] - ) - routed: Final = cast(Sequence[Mapping[str, object]], rows) # cast-ok: router-core's wildcard rows - return [cast(Mapping[str, object], row["pattern"]) for row in routed] # cast-ok: each carries its routed copy + if not isinstance(model, str): # pyright: ignore[reportUnnecessaryIsInstance] # a caller's non-string model reads the base routers + team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None + team_owned: Final = ( + team_router.get_deployments_by_pattern(model=model) if team_router is not None else () + ) + return team_owned or self._core_pattern_router.get_deployments_by_pattern(model=model) + rows: Final = self._lookup().lookup_wildcard_route( + model, team_id, [model, f"{inferred_provider(model)}/{model}"] + ) + routed: Final = cast(Sequence[Mapping[str, object]], rows) # cast-ok: router-core's wildcard rows + return [cast(Mapping[str, object], row["pattern"]) for row in routed] # cast-ok: its routed copy def _core_wildcard_routed_models(self, model: str, team_id: str | None) -> frozenset[str]: team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None @@ -1942,23 +1958,27 @@ class HostRegistry(HostModelInfo): def _core_config_deployments(self) -> Iterator[Mapping[str, object]]: """A generator, as the Python core's: the rows are read when iteration starts.""" - yield from self._rows(self._lookup().lookup_config()) + with self._deployment_index_lock: + rows: Final = self._rows(self._lookup().lookup_config()) + yield from rows # -- names ------------------------------------------------------------------------------------ def _core_get_model_names(self, team_id: str | None = None) -> list[str]: # mutable-ok: a fresh list - names, raw = self._lookup().lookup_model_names(team_id, self._list_scope(None)) - self._raise_at(raw) - return names + with self._deployment_index_lock: + names, raw = self._lookup().lookup_model_names(team_id, self._list_scope(None)) + self._raise_at(raw) + return names def _core_get_model_ids( self, model_name: str | None = None, exclude_team_models: bool = False ) -> list[str]: # mutable-ok: a fresh list - if model_name is not None and _index_key(model_name) is None: - return [] - ids, raw = self._lookup().lookup_model_ids(model_name, exclude_team_models) - self._raise_at(raw) - return ids + with self._deployment_index_lock: + if model_name is not None and _index_key(model_name) is None: + return [] + ids, raw = self._lookup().lookup_model_ids(model_name, exclude_team_models) + self._raise_at(raw) + return ids # -- resolving a name ------------------------------------------------------------------------- @@ -2017,12 +2037,6 @@ class HostRegistry(HostModelInfo): return provider_prefixed return None - def _get_deployment_by_litellm_model(self, model: object) -> list[UntypedEntry]: # mutable-ok: a fresh list - handle: Final = self._live_handle() - if handle is None or not isinstance(model, str): - return [] - return list(self._at(handle.litellm_model_positions(model))) - # -- credentials ------------------------------------------------------------------------------ @staticmethod @@ -2041,18 +2055,20 @@ class HostRegistry(HostModelInfo): def _get_model_group_deployment_usable_by_team( self, model_group_name: str, team_id: str | None ) -> Deployment | None: - if _index_key(model_group_name) is None: - return None - position, raw = self._lookup().lookup_group_usable_by_team(model_group_name, team_id) - self._raise_at(raw) - return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry + with self._deployment_index_lock: + if _index_key(model_group_name) is None: + return None + position, raw = self._lookup().lookup_group_usable_by_team(model_group_name, team_id) + self._raise_at(raw) + return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry def _get_team_public_name_deployment(self, model_id: str, team_id: str | None) -> Deployment | None: - if team_id is None or _index_key(model_id) is None: - return None - position, raw = self._lookup().lookup_team_public_name(model_id, team_id) - self._raise_at(raw) - return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry + with self._deployment_index_lock: + if team_id is None or _index_key(model_id) is None: + return None + position, raw = self._lookup().lookup_team_public_name(model_id, team_id) + self._raise_at(raw) + return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry def _get_wildcard_deployment_usable_by_team(self, model_id: str, team_id: str | None) -> Deployment | None: team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None @@ -2105,7 +2121,10 @@ class HostRegistry(HostModelInfo): # -- routable names --------------------------------------------------------------------------- def _routable_deployments(self, model_name: str, team_id: str | None) -> tuple[UntypedEntry, ...]: - return tuple(self._rows(self._lookup().lookup_routable(self._core_routable_model_group(model_name), team_id))) + with self._deployment_index_lock: + return tuple( + self._rows(self._lookup().lookup_routable(self._core_routable_model_group(model_name), team_id)) + ) def _core_team_ptu_ceiling(self, team_id: str, model_group: str) -> "PTUTeamCeiling | None": """The per-minute ceiling ``team_id``'s PTU shares set on the group serving ``model_group``, @@ -2169,24 +2188,26 @@ class HostRegistry(HostModelInfo): """The deployment ids that could serve ``model`` for ``team_id``: the alias target, then a routing group, then a team public name or wildcard route for a name that is no model group, then the model group (``get_candidate_model_ids_for_route``). Read-only.""" - from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + with self._deployment_index_lock: + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider - resolved: Final = self._core_get_model_from_alias(model=model) or model - group: Final = self.get_routing_group(resolved) if self._core_routing_groups else None - ids, raw = self._lookup().lookup_candidate_ids( - resolved, - team_id, - None if group is None else list(group.models), - [resolved, f"{inferred_provider(resolved)}/{resolved}"], - ) - self._raise_at(raw) - return frozenset(ids) + resolved: Final = self._core_get_model_from_alias(model=model) or model + group: Final = self.get_routing_group(resolved) if self._core_routing_groups else None + ids, raw = self._lookup().lookup_candidate_ids( + resolved, + team_id, + None if group is None else list(group.models), + [resolved, f"{inferred_provider(resolved)}/{resolved}"], + ) + self._raise_at(raw) + return frozenset(ids) def _core_team_model_has_alternatives(self, deployment_id: str) -> bool: """Whether a team deployment's public name has another unpaused deployment of that team.""" - found, raw = self._lookup().lookup_team_model_has_alternatives(deployment_id) - self._raise_at(raw) - return found + with self._deployment_index_lock: + found, raw = self._lookup().lookup_team_model_has_alternatives(deployment_id) + self._raise_at(raw) + return found def _core_resolve_served_model_name(self, model_name: str, served_model: str) -> str: """``served_model`` when ``model_name`` routes to it (the alias target, a routing group @@ -2315,17 +2336,18 @@ class HostRegistry(HostModelInfo): def _core_get_model_access_groups( self, model_name: str | None = None, model_access_group: str | None = None, team_id: str | None = None ) -> Mapping[str, Sequence[str]]: - use_cache: Final = model_name is None and model_access_group is None and team_id is None - if use_cache and self._access_groups_cache is not None: - return self._access_groups_cache - pairs, raw = self._lookup().lookup_access_groups( - model_name, model_access_group, team_id, self._list_scope(model_name) - ) - self._raise_at(raw) - groups: Final[dict[str, list[str]]] = {g: list(n) for g, n in pairs} # mutable-ok: as a defaultdict - if use_cache: - self._access_groups_cache = groups - return groups + with self._deployment_index_lock: + use_cache: Final = model_name is None and model_access_group is None and team_id is None + if use_cache and self._access_groups_cache is not None: + return self._access_groups_cache + pairs, raw = self._lookup().lookup_access_groups( + model_name, model_access_group, team_id, self._list_scope(model_name) + ) + self._raise_at(raw) + groups: Final[dict[str, list[str]]] = {g: list(n) for g, n in pairs} # mutable-ok: as a defaultdict + if use_cache: + self._access_groups_cache = groups + return groups def _core_is_model_access_group_for_wildcard_route(self, model_access_group: str) -> bool: groups: Final = self._core_get_model_access_groups(model_access_group=model_access_group) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3196b3c0c08..a4c16fbeaa3 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3965,6 +3965,7 @@ "supports_vision": true, "supports_output_config": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 4096 }, "azure_ai/claude-opus-4-7": { @@ -3997,6 +3998,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 2048 }, "azure_ai/claude-fable-5": { @@ -4031,6 +4033,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": false, + "supports_web_search": true, "prompt_cache_min_tokens": 512 }, "azure_ai/claude-fable-5-1": { @@ -4065,6 +4068,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": false, + "supports_web_search": true, "prompt_cache_min_tokens": 512, "deprecation_date": "2027-12-05", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement" @@ -4136,6 +4140,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 512 }, "azure_ai/claude-opus-5-5": { @@ -4171,6 +4176,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 512, "source": "https://management.azure.com/subscriptions/c873328e-b572-4770-8dff-aaeb6f1f0e79/providers/Microsoft.CognitiveServices/locations/eastus2/models?api-version=2024-10-01" }, @@ -4208,6 +4214,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 1024 }, "azure_ai/claude-sonnet-4-5": { @@ -4267,6 +4274,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 1024 }, "azure_ai/claude-sonnet-4-6": { @@ -72516,6 +72524,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.4-mini": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 8.25e-08, @@ -72537,6 +72546,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.4-nano": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 2.2e-08, @@ -73140,6 +73150,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.4-mini": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 8.25e-08, @@ -73161,6 +73172,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.4-nano": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 2.2e-08, @@ -76633,6 +76645,7 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, "supports_forced_tool_use": false, + "supports_web_search": true, "thinking_always_on": true, "prompt_cache_min_tokens": 512, "source": "https://management.azure.com/subscriptions/c873328e-b572-4770-8dff-aaeb6f1f0e79/providers/Microsoft.CognitiveServices/locations/eastus2/models?api-version=2024-10-01" @@ -78204,6 +78217,7 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, "supports_forced_tool_use": true, + "supports_web_search": true, "thinking_always_on": false, "prompt_cache_min_tokens": 512, "source": "https://platform.claude.com/docs/en/models/haiku-5-5/migration-guide" @@ -78910,6 +78924,7 @@ "prompt_cache_min_tokens": 4096.0, "rpm": 2000.0, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -78970,6 +78985,7 @@ "prompt_cache_min_tokens": 6144.0, "regional_endpoint_uplift_multiplier": 1.1, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -79029,6 +79045,7 @@ "prompt_cache_min_tokens": 6144.0, "regional_endpoint_uplift_multiplier": 1.1, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -79270,6 +79287,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "baseten/meta-models/muse-spark-open": { @@ -88127,6 +88145,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "azure_ai/claude-mythos-5-2-fast": { @@ -88162,6 +88181,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "azure_ai/claude-fable-5-2": { @@ -88193,6 +88213,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "azure_ai/claude-fable-5-2-fast": { @@ -88224,6 +88245,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "bedrock/us-gov-east-1/anthropic.claude-mythos-5-2": { diff --git a/tests/integration/management/test_budget_updates.py b/tests/integration/management/test_budget_updates.py index 6939b447866..aa351078ff9 100644 --- a/tests/integration/management/test_budget_updates.py +++ b/tests/integration/management/test_budget_updates.py @@ -36,7 +36,8 @@ def test_shortening_budget_duration_moves_reset_at_onto_the_new_schedule(gateway BUDGET_COLUMNS: Final = ( "budget_id, max_budget, soft_budget, max_parallel_requests, tpm_limit, rpm_limit, tpd_limit, model_max_budget, " - "budget_duration, allowed_models, temp_budget_increase, temp_budget_expiry, created_by, updated_by" + "model_tpm_limit, model_rpm_limit, budget_duration, allowed_models, temp_budget_increase, temp_budget_expiry, " + "created_by, updated_by" ) VOLATILE_INFO_FIELDS: Final = frozenset({"created_at", "updated_at", "budget_reset_at"}) UNLOADED_RELATIONS: Final[dict[str, JsonValue]] = dict.fromkeys( @@ -128,6 +129,8 @@ def test_partial_update_keeps_sibling_budget_fields(gateway: Gateway) -> None: "rpm_limit": 60, "tpd_limit": 100000, "model_max_budget": {"openai/gpt-4o-mini": {"max_budget": 1.0, "budget_duration": "1d"}}, + "model_tpm_limit": None, + "model_rpm_limit": None, "budget_duration": "30d", "allowed_models": [], "temp_budget_increase": None, diff --git a/tests/integration/management/test_customer_unknown_fields.py b/tests/integration/management/test_customer_unknown_fields.py new file mode 100644 index 00000000000..235d5745791 --- /dev/null +++ b/tests/integration/management/test_customer_unknown_fields.py @@ -0,0 +1,55 @@ +import uuid +from typing import Final + +import httpx +import pytest + +from tests.integration._support.client import Gateway, Scenario + +pytestmark: Final = pytest.mark.shares_worker_state + + +def _customer(scenario: Scenario, **fields: object) -> str: + customer: Final = f"integration-customer-{uuid.uuid4().hex}" + scenario.gateway.post("/customer/new", {"user_id": customer, **fields}) # pyright: ignore[reportArgumentType] # JsonValue kwargs + scenario.cleanups.callback(scenario.gateway.post, "/customer/delete", {"user_ids": [customer]}) + return customer + + +def _post(gateway: Gateway, route: str, body: dict[str, object]) -> httpx.Response: + return gateway.request("POST", route, body) # pyright: ignore[reportArgumentType] # JsonValue body + + +def _info(gateway: Gateway, customer: str) -> dict[str, object]: + return gateway.get("/customer/info", {"end_user_id": customer}) + + +@pytest.mark.parametrize( + ("body", "unknown_field"), + [ + pytest.param({"tpm_limit": 5}, "tpm_limit", id="budget_field_on_update"), + pytest.param({"quota": {"opus": 1}}, "quota", id="made_up_field"), + ], +) +def test_customer_update_rejects_unknown_fields_and_leaves_the_customer_unchanged( + gateway: Gateway, body: dict[str, object], unknown_field: str +) -> None: + with gateway.scenario() as scenario: + customer: Final = _customer(scenario, alias="before") + rejected: Final = _post(gateway, "/customer/update", {"user_id": customer, **body}) + assert rejected.status_code == 422, rejected.text + assert unknown_field in rejected.text, rejected.text + assert _info(gateway, customer)["alias"] == "before" + + accepted: Final = _post(gateway, "/customer/update", {"user_id": customer, "alias": "after"}) + assert accepted.status_code == 200, accepted.text + assert _info(gateway, customer)["alias"] == "after" + + +def test_customer_new_rejects_a_misspelled_field_and_creates_nothing(gateway: Gateway) -> None: + customer: Final = f"integration-customer-{uuid.uuid4().hex}" + rejected: Final = _post(gateway, "/customer/new", {"user_id": customer, "tpm_limti": 5}) + assert rejected.status_code == 422, rejected.text + assert "tpm_limti" in rejected.text, rejected.text + missing: Final = gateway.request("GET", "/customer/info", params={"end_user_id": customer}) + assert missing.status_code == 404, missing.text diff --git a/tests/integration/mcp/test_interactions.py b/tests/integration/mcp/test_interactions.py index 5577cf16067..b77deb5236a 100644 --- a/tests/integration/mcp/test_interactions.py +++ b/tests/integration/mcp/test_interactions.py @@ -1227,6 +1227,7 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( "PYTHONPATH": os.pathsep.join((str(source), str(source / "tests"))), "LITELLM_SALT_KEY": "subscription-pagination-key", **({"LITELLM_MCP_TOOL_LISTING_TIMEOUT": "2"} if access == "probe-timeout" else {}), + **({"LITELLM_MCP_METADATA_TIMEOUT": "5"} if access == "mixed-capabilities" else {}), "STORE_MODEL_IN_DB": "False", "DISABLE_SCHEMA_UPDATE": "true", "LITELLM_LOCAL_MODEL_COST_MAP": "True", diff --git a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py index 380881a6796..86a8cb3df4f 100644 --- a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py +++ b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py @@ -332,6 +332,22 @@ def test_utils_token_counter_local_mode_never_calls_the_peer_for_a_leading_syste assert peer.drain() == () +def test_utils_token_counter_call_endpoint_counts_a_prompt_as_one_user_message( + deployment: _Deployment, gateway: Gateway +) -> None: + with _listening(deployment) as peer: + response: Final = gateway.request( + "POST", + "/utils/token_counter", + {"model": deployment.model, "prompt": USER_TEXT}, + params={"call_endpoint": "true"}, + ) + payload: Final = _payload(response) + assert (payload["total_tokens"], payload["count_source"]) == (_COUNT, "provider"), response.text + assert payload["tokenizer_type"] == deployment.provider.tokenizer, response.text + assert _bodies(peer, deployment.provider) == ({"model": _MODEL, "messages": [USER]},) + + def test_responses_input_tokens_lifts_instructions(deployment: _Deployment, gateway: Gateway) -> None: with _listening(deployment) as peer: response: Final = gateway.request( diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index 2a0508402ae..07394c03f0c 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -550,6 +550,32 @@ def test_runtime_404_and_mantle_404_surface_provider_error(counting_proxy: Owned assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) +def test_utils_token_counter_prompt_answers_the_count_model_404_like_the_messages_form( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_rejecting(404)), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _metadata_deployment(scenario, runtime.url) + from_prompt: Final = gateway.request( + "POST", + "/utils/token_counter", + {"model": model, "prompt": "Count this message"}, + params={"call_endpoint": "true"}, + ) + from_messages: Final = gateway.request( + "POST", "/utils/token_counter", _bare(model), params={"call_endpoint": "true"} + ) + assert from_prompt.status_code == from_messages.status_code == 404, from_prompt.text + assert _REJECTION in from_prompt.text, from_prompt.text + assert _JSON_OBJECT.validate_json(from_prompt.content) == _JSON_OBJECT.validate_json(from_messages.content) + assert runtime.drain() == () + assert _mantle_bodies(mantle) == (_MANTLE_BARE, _MANTLE_BARE) + + def test_runtime_404_on_nova_does_not_call_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: gateway: Final = counting_proxy.gateway with ( diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py index a263472a45f..30a0635f3b8 100644 --- a/tests/integration/sdk/test_http2_wire.py +++ b/tests/integration/sdk/test_http2_wire.py @@ -531,6 +531,9 @@ def _serve_h2(context: ssl.SSLContext, raw: socket.socket, drain: bool) -> None: tls.sendall(connection.data_to_send()) while data := tls.recv(65535): if _reply_to_ended_streams(connection, tls, data, drain) and drain: + tls.shutdown(socket.SHUT_WR) + while tls.recv(65535): + pass return diff --git a/tests/integration/sdk/test_unset_effort_tool_bridge_wire.py b/tests/integration/sdk/test_unset_effort_tool_bridge_wire.py new file mode 100644 index 00000000000..11740f1d163 --- /dev/null +++ b/tests/integration/sdk/test_unset_effort_tool_bridge_wire.py @@ -0,0 +1,146 @@ +import json +from typing import Final + +import pytest +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm import Router +from tests.integration._support.wire import Reply, Request, wire_server + +_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} +_ARGUMENTS: Final = '{"city":"Paris"}' + + +def _tool_reply(request: Request) -> Reply: + body: Final = _OBJECT.validate_json(request.body) + if "/responses" in request.target: + return Reply( + body=json.dumps( + { + "id": "resp_weather", + "object": "response", + "created_at": 1789788253, + "model": "gpt-5.6", + "status": "completed", + "output": [ + { + "type": "function_call", + "id": "fc_weather", + "call_id": "call_weather", + "name": "get_weather", + "arguments": _ARGUMENTS, + "status": "completed", + } + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + return Reply( + body=json.dumps( + { + "id": "chatcmpl-weather", + "object": "chat.completion", + "created": 1789788253, + "model": body.get("model"), + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": _ARGUMENTS}, + } + ], + }, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + +@pytest.mark.parametrize( + ("model", "placement", "endpoint"), + ( + ("gpt-5.4-mini", "unset", "chat/completions"), + ("gpt-5.5", "unset", "chat/completions"), + ("gpt-5.6", "unset", "responses"), + ("gpt-5.4-mini", "extra-body", "responses"), + ("gpt-5.4-mini", "none", "chat/completions"), + ("gpt-5.4-mini", "deployment", "responses"), + ), +) +def test_function_tools_reach_the_chat_servable_endpoint_and_keep_tool_call_shape( + model: str, placement: str, endpoint: str +) -> None: + with wire_server(_tool_reply) as wire: + parameters: Final = { + "model": f"azure/{model}", + "api_base": wire.url, + "api_key": "synthetic-azure-key", + "api_version": "2025-04-01-preview", + } + fields: Final = ( + {"extra_body": {"reasoning_effort": "medium"}} + if placement == "extra-body" + else {"reasoning_effort": "none"} + if placement == "none" + else {} + ) + messages: Final = [{"role": "user", "content": "What's the weather in Paris? Use the tool."}] + if placement == "deployment": + router: Final = Router( + model_list=[ + { + "model_name": "weather", + "litellm_params": {**parameters, "reasoning_effort": "medium"}, + } + ], + num_retries=0, + ) + response: Final = router.completion(model="weather", messages=messages, tools=[_TOOL]) + else: + response_direct: Final = litellm.completion( + **parameters, + **fields, + messages=messages, + tools=[_TOOL], + num_retries=0, + ) + answer: Final = response if placement == "deployment" else response_direct + (sent,) = tuple(request for request in wire.drain() if request.method == "POST") + expected_path: Final = ( + "/openai/responses" if endpoint == "responses" else f"/openai/deployments/{model}/chat/completions" + ) + assert sent.target.split("?")[0] == expected_path, sent.target + assert answer.choices[0].finish_reason == "tool_calls" + assert answer.choices[0].message.tool_calls[0].function.name == "get_weather" + assert json.loads(answer.choices[0].message.tool_calls[0].function.arguments) == {"city": "Paris"} + outbound: Final = _OBJECT.validate_json(sent.body) + if endpoint == "responses": + assert outbound["tools"] == [ + { + "type": "function", + "name": "get_weather", + "parameters": _TOOL["function"]["parameters"], + "description": None, + "strict": None, + } + ] + else: + assert outbound["tools"] == [_TOOL] diff --git a/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py b/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py index bad7b3f1750..a76d5e82a37 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py @@ -759,7 +759,7 @@ def test_translate_openai_reasoning_items_are_not_duplicated_when_thinking_block def _replayed_reasoning_request(tools: list[dict[str, object]]) -> AnthropicMessagesRequest: return AnthropicMessagesRequest( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", max_tokens=64, messages=[ AnthropicMessagesUserMessageParam(role="user", content="Which city is colder?"), @@ -5820,7 +5820,7 @@ _RETRY_REPLY: Final = { "object": "response", "created_at": 1700000000, "status": "completed", - "model": "gpt-5.4-nano", + "model": "gpt-5.6", "output": [ { "type": "message", @@ -5868,7 +5868,7 @@ async def test_async_messages_retry_once_without_reasoning_the_backend_cannot_de with caplog.at_level(logging.WARNING, logger="LiteLLM"): response: Final = await litellm.anthropic.messages.acreate( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -5890,7 +5890,7 @@ async def test_async_messages_retry_once_without_reasoning_the_backend_cannot_de assert retry_warnings == ( "Backend rejected replayed encrypted reasoning, retrying once without it. " "Enable the encrypted_content_affinity pre-call check to drop such blocks before the call. " - "provider=azure model=azure/gpt-5.4-nano status=400 code=invalid_encrypted_content", + "provider=azure model=azure/gpt-5.6 status=400 code=invalid_encrypted_content", ) assert all( "gAAA_foreign" not in message and "litellm_encrypted_reasoning:" not in message for message in warnings @@ -5906,7 +5906,7 @@ def test_sync_messages_retry_once_without_reasoning_the_backend_cannot_decrypt( ) response: Final = litellm.anthropic.messages.create( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -5930,7 +5930,7 @@ async def test_async_messages_with_num_retries_send_the_rejected_reasoning_only_ ) await litellm.anthropic.messages.acreate( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -5953,7 +5953,7 @@ def test_sync_messages_with_global_num_retries_send_the_rejected_reasoning_once_ ) litellm.anthropic.messages.create( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -5981,7 +5981,7 @@ async def test_async_messages_retry_a_server_error_with_the_reasoning_and_drop_i ) response: Final = await litellm.anthropic.messages.acreate( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -6004,7 +6004,7 @@ def test_sync_messages_retry_a_server_error_with_the_reasoning_and_drop_it_only_ ) response: Final = litellm.anthropic.messages.create( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -6063,7 +6063,7 @@ async def test_async_messages_log_only_the_outcome_of_the_turn_sent_again( with pytest.raises(litellm.BadRequestError) if logged_rejections else nullcontext(): await litellm.anthropic.messages.acreate( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -6089,7 +6089,7 @@ def test_sync_messages_log_only_the_outcome_of_the_turn_sent_again( with pytest.raises(litellm.BadRequestError) if logged_rejections else nullcontext(): litellm.anthropic.messages.create( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -6124,7 +6124,7 @@ async def test_messages_do_not_retry_other_rejections( with pytest.raises(litellm.BadRequestError): await litellm.anthropic.messages.acreate( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -6277,7 +6277,7 @@ async def test_streamed_reasoning_summaries_stay_with_their_own_encrypted_item( ) stream: Final = await litellm.anthropic.messages.acreate( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index d2161fddeb0..bfbf7318d52 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -330,7 +330,7 @@ async def test_deferred_slot_keeps_the_innermost_wrapper_result(): async def test_deferred_anthropic_messages_bridged_to_the_responses_api_logs_the_provider_usage( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ): - """/v1/messages on an Azure gpt-5.4+ deployment with function tools runs three nested + """/v1/messages on an Azure deployment with function tools and explicit reasoning runs three nested wrappers: anthropic_messages, the chat adapter's acompletion, and the Responses bridge acompletion hands the call to, which retags the call as ``responses``. With logging deferred for a post-call guardrail the stored closure must carry the innermost provider @@ -382,6 +382,7 @@ async def test_deferred_anthropic_messages_bridged_to_the_responses_api_logs_the response: Final = await litellm.anthropic_messages( model="azure/gpt-5.4-nano", + reasoning_effort="medium", messages=[{"role": "user", "content": "hi"}], max_tokens=16, tools=[ diff --git a/tests/unit/proxy/management_endpoints/test_customer_budget.py b/tests/unit/proxy/management_endpoints/test_customer_budget.py index 0beca0c15e8..ac172ac2ffb 100644 --- a/tests/unit/proxy/management_endpoints/test_customer_budget.py +++ b/tests/unit/proxy/management_endpoints/test_customer_budget.py @@ -144,7 +144,6 @@ async def test_update_customer_creates_budget_with_proper_relations( update_request = UpdateCustomerRequest( user_id="test-user", max_budget=100.0, # This triggers budget creation - rpm_limit=200, # Use valid budget field ) # Act @@ -267,7 +266,6 @@ async def test_update_customer_budget_creation_with_fallback_admin( update_request = UpdateCustomerRequest( user_id="test-user", max_budget=150.0, - tpm_limit=1000, # Add another budget field ) # Act @@ -325,7 +323,6 @@ async def test_update_customer_with_budget_id_and_creation_fields( user_id="test-user", budget_id="existing-budget-link", # For linking to existing budget max_budget=300.0, # This should trigger new budget creation - rpm_limit=500, # Use valid budget field ) # Act diff --git a/tests/unit/proxy/mcp_server/test_discoverable_endpoints.py b/tests/unit/proxy/mcp_server/test_discoverable_endpoints.py index c17f4b392cd..47d24a9626f 100644 --- a/tests/unit/proxy/mcp_server/test_discoverable_endpoints.py +++ b/tests/unit/proxy/mcp_server/test_discoverable_endpoints.py @@ -16,7 +16,7 @@ from litellm.proxy._types import LiteLLM_MCPServerTable from litellm.types.mcp import MCPAuth if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Iterator, Mapping import httpx from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey @@ -15278,6 +15278,9 @@ async def test_managed_resource_metadata_excludes_refresh_permission(monkeypatch (["tools.read", "tools.write"], ("tools.write",), "present"), (None, ("tools.write",), "absent"), (["tools.read"], (), "present"), + (["channels:read,chat:write"], ("channels:read", "chat:write"), "present"), + (["channels:read,chat:write"], ("channels:read", "files:write"), "absent"), + (["channels:read"], ("channels:read", "chat:write"), "absent"), ], ) async def test_connect_requires_the_requested_vendor_grant(monkeypatch, granted, required, expected): @@ -15293,6 +15296,111 @@ async def test_connect_requires_the_requested_vendor_grant(monkeypatch, granted, read.assert_awaited_once_with(prisma, "owner", "server") +_SLACK_SCOPES: Final = frozenset(("channels:read", "chat:write", "files:write")) +_COMMA_NAMED_SCOPES: Final = frozenset(("files:read", "files:read,write")) + + +@pytest.mark.parametrize( + ("granted", "requested", "known", "expected"), + [ + (("channels:read,chat:write",), ("channels:read", "chat:write"), _SLACK_SCOPES, True), + (("channels:read,chat:write",), ("channels:read", "files:write"), _SLACK_SCOPES, False), + (("channels:read,chat:write,files:read",), ("channels:read",), _SLACK_SCOPES, False), + (("x:unknown,chat:write",), ("chat:write",), _SLACK_SCOPES, False), + (("files:read,write",), ("files:read",), _COMMA_NAMED_SCOPES, False), + (("files:read,write",), ("files:read,write",), _COMMA_NAMED_SCOPES, True), + (("channels:read,chat:write",), ("read",), _SLACK_SCOPES | {"read"}, False), + (("read", "write,admin"), ("admin",), frozenset(("read", "write", "admin")), False), + (("write,admin", "read"), ("admin",), frozenset(("read", "write", "admin")), False), + (("read", "write"), ("read", "write"), frozenset[str](), True), + ((), ("read",), frozenset(("read",)), False), + ((), (), frozenset[str](), True), + ], +) +def test_grant_covers_scopes_splits_a_lone_token_only_into_known_scopes( + granted: tuple[str, ...], requested: tuple[str, ...], known: frozenset[str], expected: bool +) -> None: + from litellm.proxy.mcp_server.oauth_utils import grant_covers_scopes + + assert grant_covers_scopes(granted, requested, lambda: known) is expected + + +@pytest.mark.parametrize( + ("granted", "requested", "expected"), + [ + (("channels:read",), ("channels:read",), True), + (("read", "write"), ("read", "write"), True), + (("read", "write"), ("admin",), False), + (("read",), ("admin",), False), + ((), (), True), + ], +) +def test_grant_covers_scopes_answers_like_a_subset_check_without_reading_server_scopes( + granted: tuple[str, ...], requested: tuple[str, ...], expected: bool +) -> None: + from litellm.proxy.mcp_server.oauth_utils import grant_covers_scopes + + def unreachable_server_scopes() -> frozenset[str]: + raise AssertionError("server scopes read for a grant with no comma token") + + assert grant_covers_scopes(granted, requested, unreachable_server_scopes) is expected + + +def test_server_scope_names_joins_configured_and_published_scopes() -> None: + from litellm.proxy.mcp_server.oauth_utils import server_scope_names + + server: Final = _cimd_oauth_server().model_copy( + update={ + "scopes": ["channels:read"], + "authorization_server_metadata": { + "issuer": "https://idp.example.com", + "authorization_endpoint": "https://idp.example.com/authorize", + "token_endpoint": "https://idp.example.com/token", + "response_types_supported": ["code"], + "scopes_supported": ["chat:write", "files:write"], + }, + } + ) + assert server_scope_names(server) == frozenset(("channels:read", "chat:write", "files:write")) + assert server_scope_names(_cimd_oauth_server()) == frozenset[str]() + assert server_scope_names(None) == frozenset[str]() + + +_PUBLISHED_SLACK_METADATA: Final = { + "issuer": "https://idp.example.com", + "authorization_endpoint": "https://idp.example.com/authorize", + "token_endpoint": "https://idp.example.com/token", + "response_types_supported": ["code"], + "scopes_supported": ["channels:read", "chat:write", "files:read"], +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "server_update", + [ + {"scopes": ["channels:read", "chat:write", "files:read"]}, + {"authorization_server_metadata": _PUBLISHED_SLACK_METADATA}, + ], +) +async def test_connect_reads_a_comma_grant_against_the_server_scopes( + monkeypatch: pytest.MonkeyPatch, server_update: "Mapping[str, object]" +) -> None: + from litellm.proxy.mcp_server import db + from litellm.proxy.mcp_server import discoverable_endpoints as endpoints + from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager + + server: Final = _cimd_oauth_server().model_copy(update=server_update) + credential: Final = {**_stored_grant(), "scopes": ["channels:read,chat:write,files:read"]} + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + monkeypatch.setattr(db, "get_user_oauth_credential", AsyncMock(return_value=credential)) + assert await endpoints._vendor_credential_state("owner", server.server_id, ("channels:read",)) == "absent" + monkeypatch.setitem(global_mcp_server_manager.registry, server.server_id, server) + assert await endpoints._vendor_credential_state("owner", server.server_id, ("channels:read",)) == "present" + assert await endpoints._vendor_credential_state("owner", server.server_id, ("files:write",)) == "absent" + assert await endpoints._vendor_credential_state("owner", server.server_id, ()) == "present" + + @pytest.mark.parametrize( ("header", "expected"), [ diff --git a/tests/unit/proxy/router_entry_point_stubs.py b/tests/unit/proxy/router_entry_point_stubs.py index 306740489d7..10176accf9f 100644 --- a/tests/unit/proxy/router_entry_point_stubs.py +++ b/tests/unit/proxy/router_entry_point_stubs.py @@ -184,6 +184,7 @@ def route_entry_points_to_stubs(router: _RouterDouble) -> _RouterDouble: _core_get_settings=lambda: router.get_settings(), _core_get_available_guardrail=lambda name: router.get_available_guardrail(guardrail_name=name), _core_get_allowed_fails_from_policy=lambda exception: router.get_allowed_fails_from_policy(exception=exception), + _core_deployment_call_timeout=lambda litellm_params: router.deployment_call_timeout(litellm_params), ) router.settings.side_effect = lambda: RouterSettings( **{name: getattr(router, name) for name in SETTING_FIELDS}, _router=settings_backend diff --git a/tests/unit/proxy/test__types.py b/tests/unit/proxy/test__types.py index 40e66f1da65..1c97e9054ae 100644 --- a/tests/unit/proxy/test__types.py +++ b/tests/unit/proxy/test__types.py @@ -26,6 +26,7 @@ from litellm.proxy._types import ( OrganizationMemberUpdateRequest, RegenerateKeyRequest, ResetSpendRequest, + UpdateCustomerRequest, UpdateKeyRequest, UpdateMCPServerRequest, UpdateTeamRequest, @@ -634,3 +635,52 @@ def test_batch_file_caps_accept_only_positive_limits(setting): with pytest.raises(ValidationError): ConfigGeneralSettings.model_validate({setting: invalid}) assert getattr(ConfigGeneralSettings.model_validate({setting: 3}), setting) == 3 + + +@pytest.mark.parametrize( + ("request_type", "body", "message"), + [ + pytest.param( + UpdateCustomerRequest, + {"user_id": "acme", "tpm_limit": 5, "quota": {"opus": 1}}, + "Unknown fields: quota, tpm_limit. tpm_limit: these live on the customer's budget object, " + "update it with /budget/update.", + id="update_with_a_budget_field_and_a_made_up_one", + ), + pytest.param( + UpdateCustomerRequest, + {"user_id": "acme", "alias": "Acme", "quota": 1}, + "Unknown fields: quota.", + id="update_with_a_made_up_field_has_no_budget_hint", + ), + pytest.param( + NewCustomerRequest, + {"user_id": "acme", "tpm_limti": 5}, + "Unknown fields: tpm_limti.", + id="create_with_a_misspelled_budget_field", + ), + ], +) +def test_a_customer_request_with_unknown_fields_is_refused_naming_them( + request_type: type[BaseModel], body: dict[str, object], message: str +) -> None: + with pytest.raises(ValidationError) as exc: + request_type.model_validate(body) + assert [error["msg"] for error in exc.value.errors()] == [f"Value error, {message}"] + + +@pytest.mark.parametrize( + ("request_type", "body"), + [ + pytest.param(UpdateCustomerRequest, {"user_id": "acme", "alias": "Acme", "blocked": True}, id="update"), + pytest.param( + NewCustomerRequest, + {"user_id": "acme", "tpm_limit": 5, "model_tpm_limit": {"opus": 60}}, + id="create_with_inline_budget_fields", + ), + ], +) +def test_a_customer_request_with_only_known_fields_is_accepted( + request_type: type[BaseModel], body: dict[str, object] +) -> None: + assert request_type.model_validate(body).model_dump(exclude_unset=True) == body diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 694b257b1f9..34868606012 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -8,10 +8,12 @@ import hashlib import itertools import json import logging +import socket import threading import traceback import uuid -from collections.abc import Callable, Coroutine, Mapping +from collections.abc import Callable, Coroutine, Iterator, Mapping +from contextlib import contextmanager from dataclasses import replace from functools import partial from pathlib import Path @@ -44,8 +46,11 @@ import litellm from litellm import Router from litellm._logging import JsonFormatter, verbose_proxy_logger from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler +from litellm.llms.azure_ai.anthropic.count_tokens import handler as azure_ai_count_tokens_http from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.bedrock.count_tokens import handler as bedrock_count_tokens_http +from litellm.llms.bedrock.count_tokens import mantle_handler as bedrock_mantle_count_tokens_http from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.kimi_k3 import ( @@ -4833,36 +4838,317 @@ async def test_local_count_of_a_prompt_with_system_or_tools_counts_them_like_the assert from_prompt.total_tokens == from_message.total_tokens > prompt_alone.total_tokens +_PROVIDER_COUNTED_DEPLOYMENTS: Final = ( + pytest.param( + {"model": "anthropic/claude-sonnet-5-5", "api_key": "test-key", "api_base": "https://count-tokens.test"}, + id="anthropic", + ), + pytest.param( + {"model": "azure_ai/claude-sonnet-5-5", "api_key": "test-key", "api_base": "https://count-tokens.test"}, + id="azure_ai", + ), + pytest.param( + {"model": "openai/gpt-5.6", "api_key": "test-key", "api_base": "https://count-tokens.test/v1"}, id="openai" + ), + pytest.param( + { + "model": "bedrock/us.anthropic.claude-sonnet-5-5", + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + }, + id="bedrock", + ), +) + + +@contextmanager +def _count_apis_answering(peer: Callable[[httpx.Request], httpx.Response]) -> Iterator[None]: + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + with ( + patch.object(anthropic_count_tokens_handler, "get_async_httpx_client", return_value=client), + patch.object(azure_ai_count_tokens_http, "get_async_httpx_client", return_value=client), + patch.object(bedrock_count_tokens_http, "get_async_httpx_client", return_value=client), + patch.object(bedrock_mantle_count_tokens_http, "get_async_httpx_client", return_value=client), + patch.object(openai_count_tokens_http, "get_async_httpx_client", return_value=client), + ): + yield + + +def _route_to_one_deployment(litellm_params: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "provider-counted", "litellm_params": litellm_params}]), + ) + + +_FOX_AS_PROMPT: Final = TokenCountRequest(model="provider-counted", prompt=_FOX) +_FOX_AS_MESSAGE: Final = TokenCountRequest(model="provider-counted", messages=[{"role": "user", "content": _FOX}]) + + @pytest.mark.asyncio -async def test_call_endpoint_prompt_keeps_the_local_count_on_a_counter_that_needs_messages( +@pytest.mark.parametrize("litellm_params", _PROVIDER_COUNTED_DEPLOYMENTS) +async def test_call_endpoint_sends_a_prompt_to_the_provider_as_one_user_message( + litellm_params: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + _route_to_one_deployment(litellm_params, monkeypatch) + sent_bodies: Final[list[JsonValue]] = [] # mutable-ok: collects what the mock peer receives + + def count_api(request: httpx.Request) -> httpx.Response: + sent_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, json={"object": "response.input_tokens", "input_tokens": 41, "inputTokens": 41}, request=request + ) + + with _count_apis_answering(count_api): + from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + from_message: Final = await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True) + + assert (from_prompt.total_tokens, from_prompt.error) == (41, False) + assert from_prompt.count_source != "local_estimate" + assert from_prompt == from_message + assert len(sent_bodies) == 2 and sent_bodies[0] == sent_bodies[1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [404, 503]) +@pytest.mark.parametrize("litellm_params", _PROVIDER_COUNTED_DEPLOYMENTS) +async def test_call_endpoint_answers_a_provider_error_on_a_prompt_like_on_the_same_text_sent_as_a_user_message( + litellm_params: dict[str, JsonValue], status: int, monkeypatch: pytest.MonkeyPatch +) -> None: + _route_to_one_deployment(litellm_params, monkeypatch) + + def count_api(request: httpx.Request) -> httpx.Response: + return httpx.Response(status, json={"error": {"message": f"scripted {status}"}}, request=request) + + with _count_apis_answering(count_api): + with pytest.raises(ProxyException) as from_prompt: + await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + with pytest.raises(ProxyException) as from_message: + await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True) + + assert from_prompt.value.code == str(status) + assert f"scripted {status}" in from_prompt.value.message + assert (from_prompt.value.code, from_prompt.value.type, from_prompt.value.message) == ( + from_message.value.code, + from_message.value.type, + from_message.value.message, + ) + + +@pytest.mark.asyncio +async def test_framed_stand_in_counts_a_prompt_like_the_same_text_sent_as_a_user_message( monkeypatch: pytest.MonkeyPatch, ) -> None: - router: Final = Router( - model_list=[ - { - "model_name": "anthropic-counted", - "litellm_params": { - "model": "anthropic/claude-sonnet-5-5", - "api_key": "test-key", - "api_base": "https://count-tokens.test", - }, - } - ] - ) - monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) - client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_unavailable_peer)) - request: Final = TokenCountRequest(model="anthropic-counted", prompt="The quick brown fox jumps over the lazy dog.") + _framed_kimi(monkeypatch) - with patch.object(anthropic_count_tokens_handler, "get_async_httpx_client", return_value=client): - with_endpoint: Final = await token_counter(request=request, call_endpoint=True) + from_prompt: Final = await token_counter(request=TokenCountRequest(model="kimi", prompt=_FOX), call_endpoint=True) + from_message: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": _FOX}]), call_endpoint=True + ) + local_prompt: Final = await token_counter(request=TokenCountRequest(model="kimi", prompt=_FOX)) + + assert (from_prompt.count_source, from_prompt.error) == ("local_estimate", False) + assert from_prompt.total_tokens == from_message.total_tokens > local_prompt.total_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", ["fireworks_ai/accounts/fireworks/models/glm-5", "deepseek/deepseek-v4-flash"], ids=["plain", "framed"] +) +async def test_call_endpoint_counts_a_prompt_on_a_provider_without_a_count_api_as_raw_text( + model: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _route_to_one_deployment({"model": model, "api_key": "test-key"}, monkeypatch) + sent: Final[list[httpx.Request]] = [] # mutable-ok: collects what the mock peer receives + + def count_api(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={"input_tokens": 41}, request=request) + + with _count_apis_answering(count_api): + from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + from_message: Final = await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True) + local_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=False) + + assert sent == [] + assert (from_prompt.count_source, from_prompt.error) == ("local_estimate", False) + assert from_prompt.total_tokens == local_prompt.total_tokens < from_message.total_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment", + [ + pytest.param( + {"litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "test-key"}}, id="cost map mode" + ), + pytest.param( + { + "litellm_params": {"model": "openai/in-house-embedder", "api_key": "test-key"}, + "model_info": {"mode": "embedding"}, + }, + id="configured mode", + ), + ], +) +async def test_call_endpoint_counts_a_prompt_on_a_non_chat_deployment_as_raw_text_without_the_provider( + deployment: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, "llm_router", Router(model_list=[{"model_name": "provider-counted", **deployment}]) + ) + sent: Final[list[httpx.Request]] = [] # mutable-ok: collects what the mock peer receives + + def count_api(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={"input_tokens": 41}, request=request) + + with _count_apis_answering(count_api): + from_endpoint: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + local: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=False) + + assert sent == [] + assert (from_endpoint.total_tokens, from_endpoint.error) == (local.total_tokens, False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["chat", "responses"]) +async def test_call_endpoint_sends_a_prompt_on_a_chat_mode_deployment_to_the_provider( + mode: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "provider-counted", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "test-key", + "api_base": "https://count-tokens.test/v1", + }, + "model_info": {"mode": mode}, + } + ] + ), + ) + sent: Final[list[httpx.Request]] = [] # mutable-ok: collects what the mock peer receives + + def count_api(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 41}, request=request) + + with _count_apis_answering(count_api): + from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + + assert (from_prompt.total_tokens, from_prompt.error) == (41, False) + assert len(sent) == 1 + + +@pytest.mark.asyncio +async def test_call_endpoint_counts_a_prompt_for_a_model_without_a_deployment_as_raw_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "other", "litellm_params": {"model": "openai/gpt-5.6"}}]), + ) + request: Final = TokenCountRequest(model="deepseek/deepseek-v4-flash", prompt=_FOX) + + from_endpoint: Final = await token_counter(request=request, call_endpoint=True) local: Final = await token_counter(request=request, call_endpoint=False) - assert (with_endpoint.error, with_endpoint.total_tokens, with_endpoint.tokenizer_type) == ( - False, - local.total_tokens, - local.tokenizer_type, + assert (from_endpoint.total_tokens, from_endpoint.error) == (local.total_tokens, False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "path"), + [("anthropic/claude-sonnet-5-5", ""), ("azure_ai/claude-sonnet-5-5", ""), ("openai/gpt-5.6", "/v1")], + ids=["anthropic", "azure_ai", "openai"], +) +@pytest.mark.parametrize( + ("row_timeout", "router_timeout"), + [({"timeout": 0.5}, None), ({"request_timeout": 0.5}, 30.0), ({}, 0.5)], + ids=["deployment timeout", "deployment request_timeout", "router_settings timeout"], +) +async def test_call_endpoint_answers_a_timeout_when_a_stalled_count_api_outlasts_the_chat_timeout( + model: str, + path: str, + row_timeout: dict[str, JsonValue], + router_timeout: float | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + with socket.create_server(("127.0.0.1", 0)) as never_answers: + port: Final = never_answers.getsockname()[1] + litellm_params: Final[dict[str, JsonValue]] = { + "model": model, + "api_key": "test-key", + "api_base": f"http://127.0.0.1:{port}{path}", + } + deployment: Final[dict[str, JsonValue]] = { + "model_name": "provider-counted", + "litellm_params": {**litellm_params, **row_timeout}, + } + monkeypatch.setattr( + litellm.proxy.proxy_server, "llm_router", Router(model_list=[deployment], timeout=router_timeout) + ) + + with pytest.raises(ProxyException) as error: + async with asyncio.timeout(10): + await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + + assert error.value.code == "408" + + +@pytest.mark.asyncio +async def test_responses_count_answers_a_timeout_when_a_stalled_count_api_outlasts_the_deployment_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with socket.create_server(("127.0.0.1", 0)) as never_answers: + port: Final = never_answers.getsockname()[1] + _route_to_one_deployment( + { + "model": "openai/gpt-5.6", + "api_key": "test-key", + "api_base": f"http://127.0.0.1:{port}/v1", + "timeout": 0.5, + }, + monkeypatch, + ) + + with pytest.raises(ProxyException) as error: + async with asyncio.timeout(10): + await count_request_tokens( + request=_FOX_AS_MESSAGE, + call_endpoint=True, + responses_request={"model": "provider-counted", "input": _FOX}, + ) + + assert error.value.code == "408" + + +@pytest.mark.asyncio +async def test_call_endpoint_names_the_status_of_a_count_api_error_with_an_empty_body( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _route_to_one_deployment( + {"model": "openai/gpt-5.6", "api_key": "test-key", "api_base": "https://count-tokens.test/v1"}, monkeypatch ) + with _count_apis_answering(lambda request: httpx.Response(404, request=request)): + with pytest.raises(ProxyException) as error: + await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + + assert error.value.code == "404" + assert "404" in error.value.message + _PINNED_COUNTS: Final = MappingProxyType({"deployment-a.test": 11, "deployment-b.test": 22}) diff --git a/tests/unit/responses/test_responses_api_bridge_flag.py b/tests/unit/responses/test_responses_api_bridge_flag.py index 210488f0888..f6084529ed6 100644 --- a/tests/unit/responses/test_responses_api_bridge_flag.py +++ b/tests/unit/responses/test_responses_api_bridge_flag.py @@ -786,17 +786,74 @@ _FUNCTION_TOOL: Final[dict[str, JsonValue]] = { @pytest.mark.parametrize("custom_llm_provider", ["openai", "azure"]) -def test_a_chat_pin_cannot_pull_gpt_5_4_function_tool_calls_off_the_responses_api(custom_llm_provider: str) -> None: - """Chat Completions rejects function tools while reasoning is on for gpt-5.4+, so the capability rule - must still bridge the request even when the deployment pins ``mode: chat``.""" +@pytest.mark.parametrize(("effort", "mode"), ((None, "chat"), ("medium", "responses"))) +def test_a_chat_pin_routes_gpt_5_4_function_tools_by_explicit_effort( + custom_llm_provider: str, effort: str | None, mode: str +) -> None: info, model = responses_api_bridge_check( model="gpt-5.4", custom_llm_provider=custom_llm_provider, tools=[_FUNCTION_TOOL], + reasoning_effort=effort, api_base="https://example.openai.azure.com" if custom_llm_provider == "azure" else None, deployment_model_info=_CHAT, ) - assert (info["mode"], model) == ("responses", "gpt-5.4") + assert (info["mode"], model) == (mode, "gpt-5.4") + + +@pytest.mark.parametrize("provider", ("openai", "azure")) +@pytest.mark.parametrize( + ("model", "effort", "extra_body", "expected"), + ( + ("gpt-5.4-mini", None, None, "chat"), + ("gpt-5.4-nano", None, None, "chat"), + ("gpt-5.5", None, None, "chat"), + ("gpt-5.6", None, None, "responses"), + ("gpt-6.1-sol", None, None, "responses"), + ("gpt-5.4-mini", "none", None, "chat"), + ("gpt-5.4-mini", "medium", None, "responses"), + ("gpt-5.4-mini", {"effort": "medium"}, None, "responses"), + ("gpt-5.4-mini", None, {"reasoning_effort": "medium"}, "responses"), + ), +) +def test_function_tools_bridge_only_when_chat_rejects_the_effective_effort( + provider: str, + model: str, + effort: str | Mapping[str, object] | None, + extra_body: Mapping[str, object] | None, + expected: str, +) -> None: + info, _ = responses_api_bridge_check( + model=model, + custom_llm_provider=provider, + tools=[_FUNCTION_TOOL], + reasoning_effort=bridge_reasoning_effort(effort, extra_body), + api_base="https://api.openai.com/v1", + ) + assert info.get("mode") == expected + + +@pytest.mark.parametrize("region", ("us", "eu")) +@pytest.mark.parametrize("model", ("gpt-5.4-mini", "gpt-5.4-nano")) +def test_regional_azure_default_none_function_tools_stay_on_chat(region: str, model: str) -> None: + info, _ = responses_api_bridge_check( + model=f"{region}/{model}", custom_llm_provider="azure", tools=[_FUNCTION_TOOL] + ) + assert info.get("mode") == "chat" + assert info.get("default_reasoning_effort") == litellm.model_cost[f"azure/{model}"]["default_reasoning_effort"] + + +def test_registered_gpt_5_6_default_none_function_tools_stay_on_chat(monkeypatch: pytest.MonkeyPatch) -> None: + model: Final = "gpt-5.6-test-default-none" + monkeypatch.setitem( + litellm.model_cost, + model, + {**litellm.model_cost["gpt-5.6"], "default_reasoning_effort": "none"}, + ) + info, _ = responses_api_bridge_check( + model=model, custom_llm_provider="openai", tools=[_FUNCTION_TOOL], api_base="https://api.openai.com/v1" + ) + assert info.get("mode") == "chat" @pytest.mark.parametrize("model", ["openai.gpt-5.6-terra", "us-east-1/openai.gpt-5.6-terra"]) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index dbd9ba8a32e..d2035b6a682 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -1418,12 +1418,7 @@ def test_responses_api_bridge_check_azure_gpt_5_4_tools_plus_reasoning_routes_to assert model_info.get("mode") == "responses" -def test_responses_api_bridge_check_azure_gpt_5_4_tools_with_default_reasoning_routes_to_responses(): - """ - Azure gpt-5.4 with tools and UNSET reasoning_effort must bridge: OpenAI enables - reasoning by default for gpt-5.4+, and Chat Completions rejects function tools - whenever reasoning is on. - """ +def test_responses_api_bridge_check_azure_gpt_5_4_tools_with_default_reasoning_stays_chat(): from litellm.main import responses_api_bridge_check with patch("litellm.responses.bridge_check.get_model_info_helper") as mock_get_model_info: @@ -1436,15 +1431,10 @@ def test_responses_api_bridge_check_azure_gpt_5_4_tools_with_default_reasoning_r ) assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" + assert model_info.get("mode") != "responses" -def test_responses_api_bridge_check_gpt_5_4_tools_with_default_reasoning_routes_to_responses(): - """ - gpt-5.4 with tools and UNSET reasoning_effort must bridge: OpenAI enables reasoning - by default for gpt-5.4+, and Chat Completions rejects function tools whenever - reasoning is on ("use /v1/responses or set reasoning_effort to 'none'"). - """ +def test_responses_api_bridge_check_gpt_5_4_tools_with_default_reasoning_stays_chat(): from litellm.main import responses_api_bridge_check with patch("litellm.responses.bridge_check.get_model_info_helper") as mock_get_model_info: @@ -1457,7 +1447,7 @@ def test_responses_api_bridge_check_gpt_5_4_tools_with_default_reasoning_routes_ ) assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" + assert model_info.get("mode") != "responses" @pytest.mark.parametrize("region", ("us", "eu")) @@ -1972,8 +1962,7 @@ def test_responses_api_bridge_check_custom_api_base_with_explicit_effort_still_r assert model_info.get("mode") == "responses" -def test_responses_api_bridge_check_azure_with_api_base_and_unset_effort_routes(): - """Azure OpenAI always sets api_base and does enforce the constraint; keep bridging.""" +def test_responses_api_bridge_check_azure_gpt_5_4_with_api_base_and_unset_effort_stays_chat(): from litellm.main import responses_api_bridge_check with patch("litellm.responses.bridge_check.get_model_info_helper") as mock_get_model_info: @@ -1987,7 +1976,7 @@ def test_responses_api_bridge_check_azure_with_api_base_and_unset_effort_routes( ) assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" + assert model_info.get("mode") != "responses" _FOUNDRY_API_BASE: Final = "https://myproject.services.ai.azure.com" diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 62c1fab244c..b270d385693 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -958,23 +958,26 @@ def test_gemini_4_argon_google_rows_carry_supports_continuation_in_both_cost_map assert _supports_continuation_rows(BACKUP_PRICES_PATH) == main_rows -def test_gemini_4_carbon_rows_share_the_gemini_3_search_allowance_in_both_cost_map_copies() -> None: +def test_gemini_4_rows_share_the_gemini_3_search_allowance_in_both_cost_map_copies() -> None: main_rows: Final = _search_free_requests_rows(PRICES_PATH) # 5,000 free grounding requests per month shared across Gemini 3 and newer: # https://ai.google.dev/gemini-api/docs/pricing#gemini-4-carbon and # https://cloud.google.com/vertex-ai/generative-ai/pricing (2026-11-10) - carbon_rows: Final = { + gemini_4_rows: Final = { + "gemini/gemini-4-argon", + "vertex_ai/gemini-4-argon", + "vertex_ai/gemini-4-argon-cyber-permissive", "gemini/gemini-4-carbon", "vertex_ai/gemini-4-carbon", "vertex_ai/gemini-4-carbon-cyber-permissive", } - assert carbon_rows <= main_rows.keys(), sorted(main_rows) + assert gemini_4_rows <= main_rows.keys(), sorted(main_rows) gemini_3_allowances: Final = { json.dumps(main_rows[key]["search_free_requests"], sort_keys=True) for key in ("gemini/gemini-3.8-flash", "vertex_ai/gemini-3.8-flash") } assert { - json.dumps(main_rows[key]["search_free_requests"], sort_keys=True) for key in carbon_rows + json.dumps(main_rows[key]["search_free_requests"], sort_keys=True) for key in gemini_4_rows } == gemini_3_allowances assert _search_free_requests_rows(BACKUP_PRICES_PATH) == main_rows @@ -1045,3 +1048,16 @@ def test_bedrock_gpt_oss_rows_declare_vision_support_like_their_safeguard_siblin def test_bedrock_gpt_oss_vision_is_explicitly_disabled_for_the_router_gate() -> None: assert litellm.utils.is_vision_explicitly_disabled("bedrock/openai.gpt-oss-120b-1:0") is True assert litellm.utils.is_vision_explicitly_disabled("bedrock_mantle/openai.gpt-oss-20b") is True + + +@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) +def test_azure_ai_claude_rows_priced_for_web_search_say_they_support_it(path: Path) -> None: + prices: Final[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]).validate_json(path.read_bytes()) + priced_rows: Final = { + key: row + for key, row in prices.items() + if key.startswith("azure_ai/claude") and isinstance(row, dict) and row.get("search_context_cost_per_query") + } + assert priced_rows + unflagged: Final = tuple(key for key, row in priced_rows.items() if row.get("supports_web_search") is not True) + assert unflagged == (), unflagged diff --git a/tests/unit/test_router/core/test_registry_characterization_lookups.py b/tests/unit/test_router/core/test_registry_characterization_lookups.py index 19d58ff10d4..94c88de73fe 100644 --- a/tests/unit/test_router/core/test_registry_characterization_lookups.py +++ b/tests/unit/test_router/core/test_registry_characterization_lookups.py @@ -4,6 +4,7 @@ access groups; blocked and unhealthy groups; credentials lookups. """ import time +from typing import Final import pytest @@ -680,3 +681,17 @@ def test_assigning_model_list_swaps_the_list_without_reindexing(make_core: MakeC found = core.registry().deployment("a") assert found is not None assert found.model_name == "n" + + +def test_the_routing_groups_source_reads_its_rows_once(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], + routing_groups=[{"group_name": "grp", "models": ["g"], "routing_strategy": "simple-shuffle"}], + ) + first: Final = core.registry().deployments(source="routing_groups") + assert ids(first) == ["a", "b"] + assert core.registry().deployments(source="routing_groups") is first + + +def test_the_credential_deployment_of_a_non_string_id_is_none(make_core: MakeCore) -> None: + assert _core(make_core).registry().credentials(123, form="deployment") is None # pyright: ignore[reportArgumentType] # a caller's non-string id diff --git a/tests/unit/test_router/core/test_registry_characterization_writes.py b/tests/unit/test_router/core/test_registry_characterization_writes.py index 0558a16fe9d..23fa2a8e0e9 100644 --- a/tests/unit/test_router/core/test_registry_characterization_writes.py +++ b/tests/unit/test_router/core/test_registry_characterization_writes.py @@ -5,7 +5,11 @@ what ``add_deployment``, ``set_model_list``, upserts and deletes leave behind in and ``registry()`` views. """ +import threading import warnings +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from typing import Final import pytest @@ -514,3 +518,80 @@ async def test_the_deprecated_arefresh_model_info_warns(make_core: MakeCore) -> await core.arefresh_model_info() # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public assert any("arefresh_model_info is deprecated" in str(w.message) for w in caught) + + +# -- reads racing writes ---------------------------------------------------------------------------- + +_GROUPS: Final = 4 + + +def _two_per_group(make_core: MakeCore) -> RouterCore: + return make_core(model_list=[deployment(f"m{index % _GROUPS}", f"d{index}") for index in range(2 * _GROUPS)]) + + +def _group_ids(core: RouterCore) -> Mapping[int, Sequence[str]]: + return {group: ids(core.registry().deployments(f"m{group}", source="all")) for group in range(_GROUPS)} + + +@pytest.mark.parametrize("upserted", [0, 3, _GROUPS + 1]) +def test_a_read_on_another_thread_waits_for_an_upsert_to_read_the_list_back(make_core: MakeCore, upserted: int) -> None: + """A lookup's positions are read against the mirror of the same registry version: a read that + starts once router-core holds a write waits until the mirror is read back, so it never reads + past the list nor, after a removal, the neighbouring group's entries.""" + core: Final = _two_per_group(make_core) + hosted: Final = core._hosted # pyright: ignore[reportAttributeAccessIssue] # the registry under test + write: Final[Callable[[Mapping[str, object]], None]] = hosted.changes_runner + started: Final = threading.Event() + pool: Final = ThreadPoolExecutor(max_workers=1) + reads: Final[list[object]] = [] # mutable-ok: one per write + + def read() -> Mapping[int, Sequence[str]]: + started.set() + return _group_ids(core) + + def write_with_a_read_started(changes: Mapping[str, object]) -> None: + write(changes) + started.clear() + reads.append(pool.submit(read)) + assert started.wait(timeout=5) + # The reader runs while router-core holds the write and the mirror does not yet. + _ = threading.Event().wait(0.02) + + hosted.changes_runner = write_with_a_read_started + try: + _ = core.apply(ChangeSet(upsert=[typed(f"m{upserted % _GROUPS}", f"d{upserted}", timeout=11)])) + finally: + hosted.changes_runner = write + + assert reads + for future in reads: + seen = future.result(timeout=5) # pyright: ignore[reportAttributeAccessIssue] # a Future + for group in range(_GROUPS): + # Between the upsert's removal and its add, its group is served by the other deployment. + assert seen[group], (group, seen) + assert set(seen[group]) <= {f"d{group}", f"d{group + _GROUPS}"}, (group, seen) + pool.shutdown() + assert sorted(_group_ids(core)[upserted % _GROUPS]) == sorted( + [f"d{upserted % _GROUPS}", f"d{upserted % _GROUPS + _GROUPS}"] + ) + + +def test_appending_to_the_live_model_list_reads_the_entry_back(make_core: MakeCore) -> None: + core: Final = _two_per_group(make_core) + live: Final = core._hosted._core_model_list # pyright: ignore[reportAttributeAccessIssue] # the live list + live.append(deployment("m0", "d8")) + assert ids(live) == [*(f"d{index}" for index in range(2 * _GROUPS)), "d8"] + + +def test_deleting_a_live_model_list_entry_reads_the_list_back(make_core: MakeCore) -> None: + core: Final = _two_per_group(make_core) + live: Final = core._hosted._core_model_list # pyright: ignore[reportAttributeAccessIssue] # the live list + del live[0] + assert ids(live) == [f"d{index}" for index in range(1, 2 * _GROUPS)] + + +def test_ending_a_build_that_never_began_leaves_the_list(make_core: MakeCore) -> None: + core: Final = _two_per_group(make_core) + live: Final = core._hosted._core_model_list # pyright: ignore[reportAttributeAccessIssue] # the live list + live.end_build() + assert ids(live) == [f"d{index}" for index in range(2 * _GROUPS)] diff --git a/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/chargebackRows.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/chargebackRows.test.ts new file mode 100644 index 00000000000..041d683803c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/chargebackRows.test.ts @@ -0,0 +1,106 @@ +import { describe, expect, test } from "vitest"; +import type { DailyData, SpendMetrics } from "@/components/UsagePage/types"; +import type { CustomerUsage } from "../customers/_components/customerPortalData"; +import { rowsByCustomer, rowsByCustomerAndKey, rowsByCustomerAndModel, sumRows, toCsv } from "./chargebackRows"; + +const totals = (spend: number, requests = 1) => ({ + spend, + requests, + inputTokens: requests * 10, + outputTokens: requests * 5, +}); + +const usage: ReadonlyMap = new Map([ + [ + "acme", + { + ...totals(6, 3), + keys: [ + { hash: "sk-hash-aaaaaaaaaaaa", alias: "prod", ...totals(5, 2) }, + { hash: "sk-hash-bbbbbbbbbbbb", alias: null, ...totals(1, 1) }, + ], + }, + ], + ["idle", { ...totals(0, 0), keys: [] }], + ["globex", { ...totals(9, 1), keys: [{ hash: "h3", alias: "main", ...totals(9, 1) }] }], +]); + +describe("rowsByCustomer", () => { + test("one row per customer with traffic, highest spend first", () => { + expect(rowsByCustomer(usage).map((r) => [r.customer, r.dimension, r.spend])).toEqual([ + ["globex", null, 9], + ["acme", null, 6], + ]); + }); +}); + +describe("rowsByCustomerAndKey", () => { + test("one row per customer and key, labelled by alias or hash prefix", () => { + expect(rowsByCustomerAndKey(usage).map((r) => [r.customer, r.dimension, r.spend])).toEqual([ + ["globex", "main", 9], + ["acme", "prod", 5], + ["acme", "sk-hash-bbbb", 1], + ]); + }); +}); + +describe("rowsByCustomerAndModel", () => { + const metrics = (spend: number): SpendMetrics => ({ + spend, + prompt_tokens: 1, + completion_tokens: 1, + total_tokens: 2, + api_requests: 1, + successful_requests: 1, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }); + const day = (models: Record): DailyData => ({ + date: "2026-11-01", + metrics: metrics(0), + breakdown: { + models: Object.fromEntries( + Object.entries(models).map(([m, s]) => [m, { metrics: metrics(s), metadata: {}, api_key_breakdown: {} }]), + ), + model_groups: {}, + mcp_servers: {}, + providers: {}, + api_keys: {}, + entities: {}, + }, + }); + test("sums each customer's models across days", () => { + const rows = rowsByCustomerAndModel(new Map([["acme", [day({ "gpt-5": 1, claude: 2 }), day({ "gpt-5": 4 })]]])); + expect(rows.map((r) => [r.customer, r.dimension, r.spend, r.requests])).toEqual([ + ["acme", "gpt-5", 5, 2], + ["acme", "claude", 2, 1], + ]); + }); +}); + +describe("sumRows and toCsv", () => { + test("total row matches the sum of the rows", () => { + const expected = { spend: 15, requests: 4, inputTokens: 40, outputTokens: 20 }; + expect(sumRows(rowsByCustomer(usage))).toEqual(expected); + }); + test("csv carries alias, optional dimension column and escaped quotes", () => { + const rows = rowsByCustomerAndKey(usage).slice(0, 1); + const csv = toCsv(rows, new Map([["globex", 'Globex "Inc"']]), "key"); + expect(csv.split("\n")).toEqual([ + '"customer_id","customer","key","requests","input_tokens","output_tokens","spend_usd"', + '"globex","Globex ""Inc""","main","1","10","5","9.000000"', + ]); + }); + test("csv neutralises ids and aliases that a spreadsheet would run as a formula", () => { + const rows = rowsByCustomer(new Map([['=HYPERLINK("x")', usage.get("acme")!]])); + const line = toCsv(rows, new Map([['=HYPERLINK("x")', "-1+1"]]), null).split("\n")[1]; + expect(line.startsWith(`"'=HYPERLINK(""x"")","'-1+1",`)).toBe(true); + expect(line.endsWith('"3","30","15","6.000000"')).toBe(true); + }); + test("csv without a dimension has no dimension column", () => { + expect(toCsv(rowsByCustomer(usage), new Map(), null).split("\n")[0]).toBe( + '"customer_id","customer","requests","input_tokens","output_tokens","spend_usd"', + ); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/chargebackRows.ts b/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/chargebackRows.ts new file mode 100644 index 00000000000..7f86b7cf6d6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/chargebackRows.ts @@ -0,0 +1,87 @@ +import type { DailyData } from "@/components/UsagePage/types"; +import { addMetrics, type CustomerUsage, type Totals } from "../customers/_components/customerPortalData"; + +export type GroupBy = "customer" | "key" | "model"; + +export interface ChargebackRow extends Totals { + customer: string; + dimension: string | null; +} + +const ZERO: Totals = { spend: 0, requests: 0, inputTokens: 0, outputTokens: 0 }; + +const hasTraffic = (r: Totals) => r.requests > 0 || r.spend > 0; +const bySpend = (a: Totals, b: Totals) => b.spend - a.spend; + +export const rowsByCustomer = (usage: ReadonlyMap): ChargebackRow[] => + [...usage.entries()] + .map(([customer, u]) => ({ + customer, + dimension: null, + spend: u.spend, + requests: u.requests, + inputTokens: u.inputTokens, + outputTokens: u.outputTokens, + })) + .filter(hasTraffic) + .sort(bySpend); + +export const rowsByCustomerAndKey = (usage: ReadonlyMap): ChargebackRow[] => + [...usage.entries()] + .flatMap(([customer, u]) => u.keys.map((k) => ({ ...k, customer, dimension: k.alias ?? k.hash.slice(0, 12) }))) + .filter(hasTraffic) + .sort(bySpend); + +export const rowsByCustomerAndModel = (perCustomer: ReadonlyMap): ChargebackRow[] => + [...perCustomer.entries()] + .flatMap(([customer, days]) => { + const byModel = new Map(); + for (const day of days) { + for (const [model, m] of Object.entries(day.breakdown.models)) + byModel.set(model, addMetrics(byModel.get(model) ?? ZERO, m.metrics)); + } + return [...byModel.entries()].map(([model, t]) => ({ ...t, customer, dimension: model })); + }) + .filter(hasTraffic) + .sort(bySpend); + +export const sumRows = (rows: readonly Totals[]): Totals => + rows.reduce( + (acc, r) => ({ + spend: acc.spend + r.spend, + requests: acc.requests + r.requests, + inputTokens: acc.inputTokens + r.inputTokens, + outputTokens: acc.outputTokens + r.outputTokens, + }), + ZERO, + ); + +const FORMULA_PREFIX = /^[=+\-@\t\r]/; +const neutralise = (v: string) => (FORMULA_PREFIX.test(v) ? `'${v}` : v); +const csvCell = (v: string | number) => `"${(typeof v === "string" ? neutralise(v) : String(v)).replace(/"/g, '""')}"`; + +export const toCsv = ( + rows: readonly ChargebackRow[], + aliasOf: ReadonlyMap, + dimension: string | null, +): string => { + const header = [ + "customer_id", + "customer", + ...(dimension ? [dimension] : []), + "requests", + "input_tokens", + "output_tokens", + "spend_usd", + ]; + const lines = rows.map((r) => [ + r.customer, + aliasOf.get(r.customer) ?? r.customer, + ...(dimension ? [r.dimension ?? ""] : []), + r.requests, + r.inputTokens, + r.outputTokens, + r.spend.toFixed(6), + ]); + return [header, ...lines].map((l) => l.map(csvCell).join(",")).join("\n"); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/page.tsx new file mode 100644 index 00000000000..262b918bde5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/page.tsx @@ -0,0 +1,175 @@ +"use client"; + +import { useMemo, useState } from "react"; +import { useTranslation } from "react-i18next"; +import { useQuery } from "@tanstack/react-query"; +import { Download } from "lucide-react"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { dailyActivityAggregatedCall } from "@/components/networking"; +import { toDailyData } from "@/components/UsagePage/dailyActivityApi"; +import { Button } from "@/components/ui/button"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { monthRange, useCustomerUsage } from "../customers/_components/customerPortalData"; +import { + rowsByCustomer, + rowsByCustomerAndKey, + rowsByCustomerAndModel, + sumRows, + toCsv, + type GroupBy, +} from "./chargebackRows"; + +const GROUPS: readonly GroupBy[] = ["customer", "key", "model"]; +const MONTHS = [0, 1, 2] as const; + +const downloadCsv = (csv: string, filename: string) => { + const url = URL.createObjectURL(new Blob([csv], { type: "text/csv" })); + Object.assign(document.createElement("a"), { href: url, download: filename }).click(); + URL.revokeObjectURL(url); +}; + +export default function ChargebacksPage() { + const { t, i18n } = useTranslation(); + const { accessToken } = useAuthorized(); + const [monthsBack, setMonthsBack] = useState(0); + const [groupBy, setGroupBy] = useState("customer"); + const customersQuery = useCustomers(100); + const customers = useMemo(() => customersQuery.data?.pages.flat() ?? [], [customersQuery.data]); + const aliasOf = useMemo(() => new Map(customers.map((c) => [c.user_id, c.alias || c.user_id])), [customers]); + const usageQuery = useCustomerUsage(accessToken, monthsBack); + const customerIds = useMemo(() => [...(usageQuery.data?.keys() ?? [])].sort(), [usageQuery.data]); + + const perModel = useQuery({ + queryKey: ["customer-portal", "usage-by-model", monthsBack, customerIds], + enabled: Boolean(accessToken) && groupBy === "model" && usageQuery.isSuccess, + queryFn: async () => { + const { start, end } = monthRange(monthsBack); + const entries = await Promise.all( + customerIds.map(async (id) => { + const request = { accessToken: accessToken!, startTime: start, endTime: end, entityIds: [id] }; + const res = await dailyActivityAggregatedCall("customer", request); + return [id, toDailyData(res)] as const; + }), + ); + return new Map(entries); + }, + }); + + const rows = useMemo(() => { + if (groupBy === "model") return perModel.data ? rowsByCustomerAndModel(perModel.data) : []; + if (!usageQuery.data) return []; + return groupBy === "key" ? rowsByCustomerAndKey(usageQuery.data) : rowsByCustomer(usageQuery.data); + }, [groupBy, usageQuery.data, perModel.data]); + const total = sumRows(rows); + const isLoading = usageQuery.isLoading || (groupBy === "model" && perModel.isLoading); + const dimensionLabel = groupBy === "customer" ? null : t(`customerPortal.chargebacks.dimension.${groupBy}`); + const monthLabel = (back: number) => + new Intl.DateTimeFormat(i18n.language, { month: "long", year: "numeric", timeZone: "UTC" }).format( + monthRange(back).start, + ); + const exportCsv = () => + downloadCsv( + toCsv(rows, aliasOf, dimensionLabel?.toLowerCase() ?? null), + `chargebacks-${monthRange(monthsBack).start.toISOString().slice(0, 7)}-${groupBy}.csv`, + ); + + return ( +
+
+
+

{t("customerPortal.chargebacks.title")}

+

{t("customerPortal.chargebacks.subtitle")}

+
+
+ + + +
+
+
+ + + + {t("customerPortal.chargebacks.columns.customer")} + {dimensionLabel && {dimensionLabel}} + {t("customerPortal.chargebacks.columns.requests")} + {t("customerPortal.chargebacks.columns.inputTokens")} + {t("customerPortal.chargebacks.columns.outputTokens")} + {t("customerPortal.chargebacks.columns.spend")} + + + + {isLoading && ( + + + {t("customerPortal.common.loading")} + + + )} + {!isLoading && rows.length === 0 && ( + + + {t("customerPortal.chargebacks.empty", { month: monthLabel(monthsBack) })} + + + )} + {rows.map((r) => ( + + +

{aliasOf.get(r.customer) ?? r.customer}

+

{r.customer}

+ + {dimensionLabel && {r.dimension}} + {t("format.number", { val: r.requests })} + {t("format.number", { val: r.inputTokens })} + {t("format.number", { val: r.outputTokens })} + {t("format.usd", { val: r.spend })} + + ))} + {rows.length > 0 && ( + + + {t("customerPortal.common.total")} + + {t("format.number", { val: total.requests })} + + {t("format.number", { val: total.inputTokens })} + + + {t("format.number", { val: total.outputTokens })} + + {t("format.usd", { val: total.spend })} + + )} + +
+
+
+ ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.tsx new file mode 100644 index 00000000000..c645c180e53 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.tsx @@ -0,0 +1,122 @@ +"use client"; + +import { useState } from "react"; +import { useTranslation } from "react-i18next"; +import { toast } from "sonner"; +import { $api } from "@/lib/http/api"; +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { buildCreateCustomerBody, digitsOnly, decimalOnly } from "./customerForms"; + +interface Props { + open: boolean; + onOpenChange: (o: boolean) => void; + onCreated: (userId: string) => void; +} + +const errorText = (e: unknown) => (e instanceof Error ? e.message : String(e)); + +export default function CreateCustomerDialog({ open, onOpenChange, onCreated }: Props) { + const { t } = useTranslation(); + const [userId, setUserId] = useState(""); + const [alias, setAlias] = useState(""); + const [budget, setBudget] = useState(""); + const [tpm, setTpm] = useState(""); + const [rpm, setRpm] = useState(""); + const create = $api.useMutation("post", "/customer/new", { + onSuccess: () => { + toast.success(t("customerPortal.toasts.created", { name: alias.trim() || userId.trim() })); + onCreated(userId.trim()); + onOpenChange(false); + }, + onError: (e: unknown) => toast.error(t("customerPortal.toasts.createFailed", { error: errorText(e) })), + }); + const submit = (e: React.FormEvent) => { + e.preventDefault(); + const fields = { userId, alias, budget, tpm, rpm }; + const body = buildCreateCustomerBody(fields); + if (body) create.mutate({ body }); + }; + return ( + + +
+ + {t("customerPortal.createDialog.title")} + {t("customerPortal.createDialog.description")} + +
+ + setUserId(e.target.value)} + placeholder={t("customerPortal.createDialog.customerIdPlaceholder")} + autoFocus + required + /> +
+
+ + setAlias(e.target.value)} + placeholder={t("customerPortal.createDialog.aliasPlaceholder")} + /> +
+
+ + setBudget(decimalOnly(e.target.value))} + placeholder={t("customerPortal.createDialog.maxBudgetPlaceholder")} + /> +
+
+
+ + setTpm(digitsOnly(e.target.value))} + placeholder={t("customerPortal.createDialog.noLimit")} + /> +
+
+ + setRpm(digitsOnly(e.target.value))} + placeholder={t("customerPortal.createDialog.noLimit")} + /> +
+
+

{t("customerPortal.createDialog.limitsHelp")}

+ + + + +
+ + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerDetail.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerDetail.tsx new file mode 100644 index 00000000000..6e9e94dedf2 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerDetail.tsx @@ -0,0 +1,156 @@ +"use client"; + +import { useState } from "react"; +import { useTranslation } from "react-i18next"; +import { useQueryClient } from "@tanstack/react-query"; +import { Ban, Pencil } from "lucide-react"; +import { toast } from "sonner"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { $api } from "@/lib/http/api"; +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { customerStatus, useCustomerUsage } from "./customerPortalData"; +import EditCustomerDialog from "./EditCustomerDialog"; +import StatusBadge from "./StatusBadge"; + +function UsageSummary({ spend, budget }: { spend: number; budget: EndUser["litellm_budget_table"] | undefined }) { + const { t } = useTranslation(); + const max = budget?.max_budget ?? null; + const tpmLimit = budget?.tpm_limit ?? null; + const rpmLimit = budget?.rpm_limit ?? null; + return ( +

+ + {t("customerPortal.detail.spendThisMonth")}{" "} + {t("format.usd", { val: spend })} + + + {max === null + ? t("customerPortal.detail.noBudget") + : t("customerPortal.detail.ofBudget", { max: t("format.usd", { val: max }) })} + + {tpmLimit !== null && ( + + {t("customerPortal.detail.tpm", { val: t("format.number", { val: tpmLimit }) })} + + )} + {rpmLimit !== null && ( + + {t("customerPortal.detail.rpm", { val: t("format.number", { val: rpmLimit }) })} + + )} +

+ ); +} + +export default function CustomerDetail({ customerId, onBack }: { customerId: string; onBack: () => void }) { + const { t } = useTranslation(); + const { accessToken } = useAuthorized(); + const info = $api.useQuery( + "get", + "/customer/info", + { params: { query: { end_user_id: customerId } } }, + { enabled: Boolean(accessToken) }, + ); + const usageQuery = useCustomerUsage(accessToken); + const [editOpen, setEditOpen] = useState(false); + const queryClient = useQueryClient(); + const refresh = () => { + info.refetch(); + queryClient.invalidateQueries({ queryKey: ["get", "/customer/list"] }); + }; + const block = $api.useMutation("post", "/customer/block", { + onSuccess: () => { + toast.success(t("customerPortal.toasts.blocked")); + refresh(); + }, + }); + const unblock = $api.useMutation("post", "/customer/update", { + onSuccess: () => { + toast.success(t("customerPortal.toasts.unblocked")); + refresh(); + }, + }); + + const c = info.data; + const usage = usageQuery.data?.get(customerId); + const spend = usage?.spend ?? 0; + const keys = usage?.keys ?? []; + const label = c?.alias || customerId; + + return ( +
+
+
+
+

{label}

+ {c && } + {customerId} +
+ +
+
+ + {c?.blocked ? ( + + ) : ( + + )} +
+
+

{t("customerPortal.common.blockHint")}

+ +
+
+

{t("customerPortal.detail.keysTitle")}

+

{t("customerPortal.detail.keysSubtitle")}

+
+ {keys.length === 0 ? ( +

{t("customerPortal.detail.keysEmpty")}

+ ) : ( + + + + {t("customerPortal.detail.columns.key")} + {t("customerPortal.detail.columns.requests")} + {t("customerPortal.detail.columns.tokens")} + {t("customerPortal.detail.columns.spend")} + + + + {keys.map((k) => ( + + +

{k.alias ?? t("customerPortal.detail.unnamedKey")}

+

{k.hash.slice(0, 12)}

+ + {t("format.number", { val: k.requests })} + + {t("format.number", { val: k.inputTokens + k.outputTokens })} + + {t("format.usd", { val: k.spend })} + + ))} + +
+ )} +
+ + {c && } +
+ ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/EditCustomerDialog.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/EditCustomerDialog.tsx new file mode 100644 index 00000000000..aebc8d60053 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/EditCustomerDialog.tsx @@ -0,0 +1,96 @@ +"use client"; + +import { useState } from "react"; +import { useTranslation } from "react-i18next"; +import { toast } from "sonner"; +import { $api } from "@/lib/http/api"; +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { buildUpdateCustomerBody, decimalOnly } from "./customerForms"; + +interface Props { + open: boolean; + onOpenChange: (o: boolean) => void; + customer: EndUser; + onSaved: () => void; +} + +const asText = (n: number | null | undefined) => (n === null || n === undefined ? "" : String(n)); + +function EditCustomerForm({ onOpenChange, customer, onSaved }: Omit) { + const { t } = useTranslation(); + const [alias, setAlias] = useState(customer.alias ?? ""); + const [budget, setBudget] = useState(asText(customer.litellm_budget_table?.max_budget)); + const budgetId = customer.litellm_budget_table?.budget_id ?? null; + const update = $api.useMutation("post", "/customer/update", { + onSuccess: () => { + toast.success(t("customerPortal.toasts.updated")); + onSaved(); + onOpenChange(false); + }, + onError: (e: unknown) => + toast.error(t("customerPortal.toasts.updateFailed", { error: e instanceof Error ? e.message : String(e) })), + }); + const submit = (e: React.FormEvent) => { + e.preventDefault(); + const fields = { userId: customer.user_id, alias, budget, blocked: Boolean(customer.blocked) }; + update.mutate({ body: buildUpdateCustomerBody(fields) }); + }; + return ( + +
+ + {t("customerPortal.editDialog.title")} + + {customer.user_id} + + +
+ + setAlias(e.target.value)} autoFocus /> +
+
+ + setBudget(decimalOnly(e.target.value))} + placeholder={t("customerPortal.editDialog.maxBudgetPlaceholder")} + /> +
+ {budgetId && ( +

+ {t("customerPortal.editDialog.sharedBudgetWarning", { budgetId })} +

+ )} + + + + +
+ + ); +} + +export default function EditCustomerDialog({ open, ...form }: Props) { + return ( + + {open && } + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/StatusBadge.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/StatusBadge.tsx new file mode 100644 index 00000000000..005171abc53 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/StatusBadge.tsx @@ -0,0 +1,15 @@ +import { useTranslation } from "react-i18next"; +import { Badge } from "@/components/ui/badge"; +import type { CustomerStatus } from "./customerPortalData"; + +const VARIANT = { active: "secondary", over_budget: "destructive", blocked: "outline" } as const; +const LABEL_KEY = { + active: "customerPortal.status.active", + over_budget: "customerPortal.status.overBudget", + blocked: "customerPortal.status.blocked", +} as const; + +export default function StatusBadge({ status }: { status: CustomerStatus }) { + const { t } = useTranslation(); + return {t(LABEL_KEY[status])}; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerForms.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerForms.test.ts new file mode 100644 index 00000000000..3c7ecc97484 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerForms.test.ts @@ -0,0 +1,49 @@ +import { describe, expect, test } from "vitest"; +import { buildCreateCustomerBody, buildUpdateCustomerBody } from "./customerForms"; + +describe("buildCreateCustomerBody", () => { + const fields = { userId: " acme ", alias: " Acme Corp ", budget: "25", tpm: "1000", rpm: "60" }; + const emptyFields = { userId: "acme", alias: "", budget: "", tpm: "", rpm: "" }; + + test("trims ids and aliases and sends numeric limits with a monthly budget reset", () => { + const expected = { + user_id: "acme", + alias: "Acme Corp", + blocked: false, + max_budget: 25, + budget_duration: "1mo", + tpm_limit: 1000, + rpm_limit: 60, + }; + expect(buildCreateCustomerBody(fields)).toEqual(expected); + }); + + test("empty optional fields become null and no budget means no budget_duration", () => { + const expected = { + user_id: "acme", + alias: null, + blocked: false, + max_budget: null, + budget_duration: null, + tpm_limit: null, + rpm_limit: null, + }; + expect(buildCreateCustomerBody(emptyFields)).toEqual(expected); + }); + + test("a blank customer id produces no request", () => { + expect(buildCreateCustomerBody({ ...fields, userId: " " })).toBeNull(); + }); +}); + +describe("buildUpdateCustomerBody", () => { + test("preserves the blocked flag and clears the budget when emptied", () => { + const fields = { userId: "acme", alias: "", budget: "", blocked: true }; + const expected = { user_id: "acme", alias: null, blocked: true, max_budget: null }; + expect(buildUpdateCustomerBody(fields)).toEqual(expected); + }); + test("parses a decimal budget", () => { + const fields = { userId: "acme", alias: "Acme", budget: "12.50", blocked: false }; + expect(buildUpdateCustomerBody(fields).max_budget).toBe(12.5); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerForms.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerForms.ts new file mode 100644 index 00000000000..f1fafc0c718 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerForms.ts @@ -0,0 +1,46 @@ +import type { components } from "@/lib/http/schema"; + +export type NewCustomerBody = components["schemas"]["NewCustomerRequest"]; +export type UpdateCustomerBody = components["schemas"]["UpdateCustomerRequest"]; + +export const digitsOnly = (s: string) => s.replace(/[^0-9]/g, ""); +export const decimalOnly = (s: string) => s.replace(/[^0-9.]/g, ""); + +const numberOrNull = (s: string): number | null => (s.trim() === "" ? null : Number(s)); + +export interface CreateCustomerFields { + userId: string; + alias: string; + budget: string; + tpm: string; + rpm: string; +} + +export const buildCreateCustomerBody = (f: CreateCustomerFields): NewCustomerBody | null => { + const userId = f.userId.trim(); + if (!userId) return null; + const maxBudget = numberOrNull(f.budget); + return { + user_id: userId, + alias: f.alias.trim() || null, + blocked: false, + max_budget: maxBudget, + budget_duration: maxBudget === null ? null : "1mo", + tpm_limit: numberOrNull(f.tpm), + rpm_limit: numberOrNull(f.rpm), + }; +}; + +export interface EditCustomerFields { + userId: string; + alias: string; + budget: string; + blocked: boolean; +} + +export const buildUpdateCustomerBody = (f: EditCustomerFields): UpdateCustomerBody => ({ + user_id: f.userId, + alias: f.alias.trim() || null, + blocked: f.blocked, + max_budget: numberOrNull(f.budget), +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPortalData.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPortalData.test.ts new file mode 100644 index 00000000000..cf48c81db5b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPortalData.test.ts @@ -0,0 +1,110 @@ +import { describe, expect, test } from "vitest"; +import type { DailyData, KeyMetricWithMetadata, MetricWithMetadata, SpendMetrics } from "@/components/UsagePage/types"; +import { customerStatus, monthRange, rankBySpend, usageByCustomer } from "./customerPortalData"; + +const metrics = (spend: number, requests = 1, prompt = 10, completion = 5): SpendMetrics => ({ + spend, + prompt_tokens: prompt, + completion_tokens: completion, + total_tokens: prompt + completion, + api_requests: requests, + successful_requests: requests, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, +}); + +const key = (alias: string | null, spend: number, requests = 1): KeyMetricWithMetadata => ({ + metrics: metrics(spend, requests), + metadata: { key_alias: alias, team_id: null }, +}); + +const entity = (spend: number, keys: Record, requests = 1): MetricWithMetadata => ({ + metrics: metrics(spend, requests), + metadata: {}, + api_key_breakdown: keys, +}); + +const day = (date: string, entities: Record): DailyData => ({ + date, + metrics: metrics(0, 0, 0, 0), + breakdown: { models: {}, model_groups: {}, mcp_servers: {}, providers: {}, api_keys: {}, entities }, +}); + +describe("usageByCustomer", () => { + test("sums a customer's spend, requests and tokens across days", () => { + const usage = usageByCustomer([ + day("2026-11-01", { acme: entity(1.5, {}, 2) }), + day("2026-11-02", { acme: entity(2.25, {}, 3) }), + ]); + const expected = { spend: 3.75, requests: 5, inputTokens: 20, outputTokens: 10 }; + expect(usage.get("acme")).toMatchObject(expected); + }); + + test("keeps customers separate", () => { + const usage = usageByCustomer([day("2026-11-01", { acme: entity(1, {}), globex: entity(4, {}) })]); + expect(usage.get("acme")?.spend).toBe(1); + expect(usage.get("globex")?.spend).toBe(4); + }); + + test("merges the same key across days and sorts keys by spend", () => { + const usage = usageByCustomer([ + day("2026-11-01", { acme: entity(3, { h1: key("prod", 1), h2: key("staging", 2) }) }), + day("2026-11-02", { acme: entity(5, { h1: key("prod", 5) }) }), + ]); + const prod = { hash: "h1", alias: "prod", spend: 6, requests: 2 }; + const staging = { hash: "h2", alias: "staging", spend: 2, requests: 1 }; + expect(usage.get("acme")?.keys).toEqual([expect.objectContaining(prod), expect.objectContaining(staging)]); + }); + + test("keeps the first non-null alias seen for a key", () => { + const usage = usageByCustomer([ + day("2026-11-01", { acme: entity(1, { h1: key(null, 1) }) }), + day("2026-11-02", { acme: entity(1, { h1: key("prod", 1) }) }), + ]); + expect(usage.get("acme")?.keys[0]?.alias).toBe("prod"); + }); +}); + +describe("customerStatus", () => { + test("blocked wins over budget", () => { + expect(customerStatus({ blocked: true, litellm_budget_table: { max_budget: 1 } }, 5)).toBe("blocked"); + }); + test("over budget once spend reaches max_budget", () => { + expect(customerStatus({ blocked: false, litellm_budget_table: { max_budget: 10 } }, 10)).toBe("over_budget"); + expect(customerStatus({ blocked: false, litellm_budget_table: { max_budget: 10 } }, 9.99)).toBe("active"); + }); + test("no budget is active regardless of spend", () => { + expect(customerStatus({ blocked: false, litellm_budget_table: null }, 1000)).toBe("active"); + }); +}); + +describe("rankBySpend", () => { + const customer = (user_id: string) => ({ user_id, blocked: false, spend: 0, litellm_budget_table: null }); + test("orders customers by this month's spend, highest first, zero for customers without traffic", () => { + const usage = usageByCustomer([day("2026-11-01", { b: entity(2, {}), c: entity(9, { h1: key("k", 9) }) })]); + const rows = rankBySpend([customer("a"), customer("b"), customer("c")], usage); + expect(rows.map((r) => [r.c.user_id, r.spend, r.keys])).toEqual([ + ["c", 9, 1], + ["b", 2, 0], + ["a", 0, 0], + ]); + }); +}); + +describe("monthRange", () => { + test("covers the whole UTC month, monthsBack months ago", () => { + const now = new Date(Date.UTC(2026, 10, 9, 12)); + expect(monthRange(0, now)).toEqual({ + start: new Date(Date.UTC(2026, 10, 1)), + end: new Date(Date.UTC(2026, 10, 30, 23, 59, 59)), + }); + expect(monthRange(1, now)).toEqual({ + start: new Date(Date.UTC(2026, 9, 1)), + end: new Date(Date.UTC(2026, 9, 31, 23, 59, 59)), + }); + }); + test("crosses a year boundary", () => { + expect(monthRange(2, new Date(Date.UTC(2026, 0, 15))).start).toEqual(new Date(Date.UTC(2025, 10, 1))); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPortalData.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPortalData.ts new file mode 100644 index 00000000000..5bc1c2ec437 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPortalData.ts @@ -0,0 +1,89 @@ +import { useQuery } from "@tanstack/react-query"; +import { dailyActivityAggregatedCall } from "@/components/networking"; +import { toDailyData } from "@/components/UsagePage/dailyActivityApi"; +import type { DailyData, SpendMetrics } from "@/components/UsagePage/types"; +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; + +export type CustomerStatus = "active" | "over_budget" | "blocked"; + +export const monthRange = (monthsBack = 0, now = new Date()): { start: Date; end: Date } => ({ + start: new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth() - monthsBack, 1)), + end: new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth() - monthsBack + 1, 0, 23, 59, 59)), +}); + +export interface Totals { + spend: number; + requests: number; + inputTokens: number; + outputTokens: number; +} + +export interface CustomerKeyUsage extends Totals { + hash: string; + alias: string | null; +} + +export interface CustomerUsage extends Totals { + keys: readonly CustomerKeyUsage[]; +} + +const ZERO: Totals = { spend: 0, requests: 0, inputTokens: 0, outputTokens: 0 }; + +export const addMetrics = (t: Totals, m: Partial | undefined): Totals => ({ + spend: t.spend + (m?.spend ?? 0), + requests: t.requests + (m?.api_requests ?? 0), + inputTokens: t.inputTokens + (m?.prompt_tokens ?? 0), + outputTokens: t.outputTokens + (m?.completion_tokens ?? 0), +}); + +export const usageByCustomer = (days: readonly DailyData[]): Map => { + const totals = new Map(); + const keys = new Map>(); + for (const day of days) { + for (const [id, entity] of Object.entries(day.breakdown.entities)) { + totals.set(id, addMetrics(totals.get(id) ?? ZERO, entity.metrics)); + const byKey = keys.get(id) ?? new Map(); + for (const [hash, key] of Object.entries(entity.api_key_breakdown)) { + const prev = byKey.get(hash) ?? { ...ZERO, hash, alias: key.metadata.key_alias }; + byKey.set(hash, { ...prev, ...addMetrics(prev, key.metrics), alias: prev.alias ?? key.metadata.key_alias }); + } + keys.set(id, byKey); + } + } + return new Map( + [...totals.entries()].map(([id, t]) => [ + id, + { ...t, keys: [...(keys.get(id)?.values() ?? [])].sort((a, b) => b.spend - a.spend) }, + ]), + ); +}; + +export const useCustomerUsage = (accessToken: string | null, monthsBack = 0) => + useQuery({ + queryKey: ["customer-portal", "usage", monthsBack], + enabled: Boolean(accessToken), + queryFn: async () => { + const { start, end } = monthRange(monthsBack); + const res = await dailyActivityAggregatedCall("customer", { + accessToken: accessToken!, + startTime: start, + endTime: end, + }); + return usageByCustomer(toDailyData(res)); + }, + }); + +export const customerStatus = (c: Pick, spend: number): CustomerStatus => { + if (c.blocked) return "blocked"; + const max = c.litellm_budget_table?.max_budget; + return max !== null && max !== undefined && spend >= max ? "over_budget" : "active"; +}; + +export const rankBySpend = (customers: readonly EndUser[], usage: ReadonlyMap | undefined) => + customers + .map((c) => { + const u = usage?.get(c.user_id); + const spend = u?.spend ?? 0; + return { c, spend, keys: u?.keys.length ?? 0, status: customerStatus(c, spend) }; + }) + .sort((a, b) => b.spend - a.spend); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/page.tsx new file mode 100644 index 00000000000..f80d15f5acf --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/page.tsx @@ -0,0 +1,153 @@ +"use client"; + +import { useMemo, useState } from "react"; +import { useTranslation } from "react-i18next"; +import { parseAsString, useQueryState } from "nuqs"; +import { Plus } from "lucide-react"; +import { toast } from "sonner"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { useCustomers, type EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { $api } from "@/lib/http/api"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import CustomerDetail from "./_components/CustomerDetail"; +import CreateCustomerDialog from "./_components/CreateCustomerDialog"; +import StatusBadge from "./_components/StatusBadge"; +import { rankBySpend, useCustomerUsage } from "./_components/customerPortalData"; + +export default function CustomersPage() { + const { t } = useTranslation(); + const { accessToken } = useAuthorized(); + const [selectedId, setSelectedId] = useQueryState("id", parseAsString); + const [search, setSearch] = useState(""); + const [createOpen, setCreateOpen] = useState(false); + const customersQuery = useCustomers(100, search || undefined); + const usageQuery = useCustomerUsage(accessToken); + const customers: EndUser[] = useMemo(() => customersQuery.data?.pages.flat() ?? [], [customersQuery.data]); + const rows = useMemo(() => rankBySpend(customers, usageQuery.data), [customers, usageQuery.data]); + const block = $api.useMutation("post", "/customer/block", { + onSuccess: () => { + toast.success(t("customerPortal.toasts.blocked")); + customersQuery.refetch(); + }, + }); + const unblock = $api.useMutation("post", "/customer/update", { + onSuccess: () => { + toast.success(t("customerPortal.toasts.unblocked")); + customersQuery.refetch(); + }, + }); + + if (selectedId) return setSelectedId(null)} />; + + return ( +
+
+
+

{t("customerPortal.list.title")}

+

{t("customerPortal.list.subtitle")}

+
+ +
+

{t("customerPortal.common.blockHint")}

+
+
+ setSearch(e.target.value)} + /> +
+ + + + {t("customerPortal.list.columns.customer")} + {t("customerPortal.list.columns.spend")} + {t("customerPortal.list.columns.budget")} + {t("customerPortal.list.columns.keys")} + {t("customerPortal.list.columns.status")} + + + + + {customersQuery.isLoading && ( + + + {t("customerPortal.list.loading")} + + + )} + {!customersQuery.isLoading && rows.length === 0 && ( + + + {t("customerPortal.list.empty")} + + + )} + {rows.map(({ c, spend, keys, status }) => { + const max = c.litellm_budget_table?.max_budget ?? null; + return ( + setSelectedId(c.user_id, { history: "push" })} + > + +

{c.alias || c.user_id}

+

{c.user_id}

+ + + {t("format.usd", { val: spend })} + + + {max === null + ? t("customerPortal.common.noBudget") + : t("customerPortal.list.budgetUsage", { + spend: t("format.usd", { val: spend }), + max: t("format.usd", { val: max }), + })} + + {keys} + + + + e.stopPropagation()}> + {c.blocked ? ( + + ) : ( + + )} + + + ); + })} + +
+
+ { + customersQuery.refetch(); + setSelectedId(id, { history: "push" }); + }} + /> +
+ ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 5c93d214812..d898f1b0373 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -116,13 +116,19 @@ export function AgentControlPlaneView() { } const FULL_BLEED_SEGMENTS = new Set(["logs"]); +const CUSTOMER_PORTAL_SEGMENTS = new Set(["customers", "chargebacks"]); function DashboardShell({ children }: { children: React.ReactNode }) { const { t } = useTranslation(); const { accessToken } = useAuth(); - const { mode } = usePluginMode(); + const { mode, setMode } = usePluginMode(); const pathname = usePathname(); const routeSegment = routeSegmentForPathname(pathname); + const isPortalRoute = CUSTOMER_PORTAL_SEGMENTS.has(routeSegment); + useEffect(() => { + if (isPortalRoute && mode !== "customer-portal") setMode("customer-portal"); + if (!isPortalRoute && mode === "customer-portal") setMode("ai-gateway"); + }, [isPortalRoute, mode, setMode]); useEffect(() => { void recordUiEvent({ page: routeSegment || "home", action: "view" }); }, [routeSegment]); @@ -142,7 +148,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) { const sidebarCollapsed = sidebarOverride?.segment === routeSegment ? sidebarOverride.collapsed : isFullBleed; const toggleSidebar = () => setSidebarOverride({ segment: routeSegment, collapsed: !sidebarCollapsed }); - const isGateway = mode === "ai-gateway"; + const isGateway = mode === "ai-gateway" || mode === "customer-portal"; // Non-gateway (agent control plane) mode keeps the original full-width Navbar, // which carries the account menu; the redesigned sidebar + header shell is diff --git a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts index 19dac805948..90b49e56c5a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts @@ -12,6 +12,8 @@ const LEGACY_PAGE_ROUTES: ReadonlyMap = new Map( chat: "chat", "access-groups": "access-groups", budgets: "budgets", + customers: "customers", + chargebacks: "chargebacks", workflows: "workflows", "guardrails-monitor": "guardrails-monitor", "mcp-servers": "mcp-servers", diff --git a/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.integration.test.tsx b/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.integration.test.tsx index 64bb8f15585..a0f6e59d029 100644 --- a/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.integration.test.tsx @@ -8,13 +8,14 @@ import { setServerRootPath } from "@/lib/serverRootPath"; import ViewSwitcher from "./ViewSwitcher"; -const { state, getUiSettingsMock, getUiConfigMock } = vi.hoisted(() => ({ +const { state, getUiSettingsMock, getUiConfigMock, pushMock } = vi.hoisted(() => ({ state: { pathname: "/ui/", plugins: [] as Plugin[] }, getUiSettingsMock: vi.fn(), getUiConfigMock: vi.fn(), + pushMock: vi.fn(), })); -vi.mock("next/navigation", () => ({ usePathname: () => state.pathname })); +vi.mock("next/navigation", () => ({ usePathname: () => state.pathname, useRouter: () => ({ push: pushMock }) })); vi.mock("@/components/networking", () => ({ getProxyBaseUrl: () => "", getUiSettings: getUiSettingsMock, @@ -56,6 +57,8 @@ describe("ViewSwitcher", () => { }); afterEach(() => { + pushMock.mockClear(); + assignSpy.mockClear(); Object.defineProperty(window, "location", { configurable: true, value: originalLocation }); testQueryClient.clear(); localStorage.clear(); @@ -103,6 +106,27 @@ describe("ViewSwitcher", () => { expect(assignSpy).not.toHaveBeenCalled(); }); + it("should open the Customer Portal at the customers page and come back to the gateway", async () => { + const user = userEvent.setup(); + renderSwitcher(); + + await user.click(screen.getByRole("button", { name: "Switch workspace, current: AI Gateway" })); + await user.click(await screen.findByRole("menuitem", { name: /Customer Portal/ })); + expect(pushMock).toHaveBeenCalledWith("/ui/customers"); + expect(assignSpy).not.toHaveBeenCalled(); + expect(localStorage.getItem("litellm_plugin_mode")).toBeNull(); + }); + + it("should route back to the gateway from Customer Portal mode instead of only flipping the mode", async () => { + localStorage.setItem("litellm_plugin_mode", "customer-portal"); + const user = userEvent.setup(); + renderSwitcher(); + + await user.click(await screen.findByRole("button", { name: "Switch workspace, current: Customer Portal" })); + await user.click(await screen.findByRole("menuitem", { name: /^AI Gateway/ })); + expect(pushMock).toHaveBeenCalledWith("/ui/api-keys"); + }); + it("should navigate from the gateway to enabled Chat without changing plugin mode", async () => { getUiSettingsMock.mockResolvedValue({ values: { enable_chat_ui: true } }); const user = userEvent.setup(); diff --git a/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.tsx b/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.tsx index 5c8d4a7a850..38b051515fb 100644 --- a/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.tsx @@ -1,7 +1,7 @@ import { useTerminology } from "@/contexts/TerminologyContext"; import React from "react"; import { useTranslation } from "react-i18next"; -import { usePathname } from "next/navigation"; +import { usePathname, useRouter } from "next/navigation"; import { Button } from "@/components/ui/button"; import { DropdownMenu, @@ -11,6 +11,7 @@ import { } from "@/components/ui/dropdown-menu"; import { Check, ChevronDown, ChevronsUpDown } from "lucide-react"; import { useSidebar } from "@/components/shared/Sidebar"; +import BetaBadge from "@/components/BetaBadge"; import { usePluginMode } from "@/contexts/PluginModeContext"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import { uiHref } from "@/utils/uiHref"; @@ -19,6 +20,7 @@ import i18n from "@/i18n/i18n"; import { productNames } from "@/i18n/untranslated/productNames"; const GATEWAY = "ai-gateway"; +const CUSTOMER_PORTAL = "customer-portal"; const CHAT = "chat"; const MOYAI = "moyai"; @@ -51,6 +53,9 @@ interface ViewSwitcherProps { contentClassName?: string; } +const pluginLabel = (mode: string, plugins: readonly { name: string; display_name: string }[], portalLabel: string) => + mode === CUSTOMER_PORTAL ? portalLabel : plugins.find((p) => p.name === mode)?.display_name; + export default function ViewSwitcher({ className, contentClassName }: ViewSwitcherProps) { const { t } = useTranslation(); const term = useTerminology(); @@ -58,6 +63,7 @@ export default function ViewSwitcher({ className, contentClassName }: ViewSwitch const { mode, setMode, plugins } = usePluginMode(); const { data: uiSettings } = useUISettings(); const pathname = usePathname(); + const router = useRouter(); const chatEnabled = Boolean(uiSettings?.values?.enable_chat_ui); const moyaiUrl = (uiSettings?.values?.moyai_url as string | undefined) ?? null; @@ -71,7 +77,7 @@ export default function ViewSwitcher({ className, contentClassName }: ViewSwitch const activeLabel = activeWorkspaceLabel( isChatRoute, isMoyaiRoute, - plugins.find((p) => p.name === mode)?.display_name, + pluginLabel(mode, plugins, t("customerPortal.switcher.label")), gatewayLabel, ); @@ -81,6 +87,14 @@ export default function ViewSwitcher({ className, contentClassName }: ViewSwitch ]; const selectMode = (key: string) => { + if (key === CUSTOMER_PORTAL) { + router.push(uiHref("customers")); + return; + } + if (mode === CUSTOMER_PORTAL && key === GATEWAY) { + router.push(uiHref("api-keys")); + return; + } setMode(key); if (isStandaloneRoute) { window.location.assign(uiHref("")); @@ -107,6 +121,21 @@ export default function ViewSwitcher({ className, contentClassName }: ViewSwitch onClick: chatEnabled ? () => window.location.assign(uiHref(CHAT)) : undefined, }; + const customerPortalItem: ViewSwitcherItem = { + key: CUSTOMER_PORTAL, + label: ( +
+ + + {t("customerPortal.switcher.label")} + {" "} + {t("customerPortal.switcher.description")} + + {!isStandaloneRoute && mode === CUSTOMER_PORTAL && } +
+ ), + onClick: () => selectMode(CUSTOMER_PORTAL), + }; const items: ViewSwitcherItem[] = [ ...modeEntries.map((e) => ({ key: e.key, @@ -118,6 +147,7 @@ export default function ViewSwitcher({ className, contentClassName }: ViewSwitch ), onClick: () => selectMode(e.key), })), + customerPortalItem, chatItem, { key: MOYAI, @@ -161,6 +191,12 @@ export default function ViewSwitcher({ className, contentClassName }: ViewSwitch <> {activeLabel} + {!isStandaloneRoute && mode === CUSTOMER_PORTAL && ( + <> + {" "} + + + )} diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index f48271a872c..6040da32f6c 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -55,12 +55,14 @@ import { Trophy, User, Users, + Receipt, Wallet, Wrench, Workflow, } from "lucide-react"; import Link from "next/link"; import ViewSwitcher from "@/components/Navbar/ViewSwitcher"; +import { usePluginMode } from "@/contexts/PluginModeContext"; import { AriaSidebarLink } from "@/components/liteadmin/LiteAdmin"; import { usePathname } from "next/navigation"; import { useMemo, useState } from "react"; @@ -197,6 +199,28 @@ export const visibleMenuGroups = (groups: readonly MenuGroup[], context: MenuVis // Menu groups organized by category - defined outside component for export. // Shape (key/page/label/roles/children) is consumed by page_utils.ts; only the // icons changed to lucide as part of the sidebar redesign. +const customerPortalGroups: MenuGroup[] = [ + { + groupLabel: "CUSTOMER PORTAL", + items: [ + { + key: "customers", + page: "customers", + label: "shell.nav.items.customers", + icon: , + roles: all_admin_roles, + }, + { + key: "chargebacks", + page: "chargebacks", + label: "shell.nav.items.chargebacks", + icon: , + roles: all_admin_roles, + }, + ], + }, +]; + const menuGroups: MenuGroup[] = [ { groupLabel: "AI GATEWAY", @@ -542,6 +566,7 @@ const findMenuItemKey = (route: string): string => { const GROUP_LABELS: Record = { "AI GATEWAY": "shell.nav.groups.aiGateway", + "CUSTOMER PORTAL": "shell.nav.groups.customerPortal", OBSERVABILITY: "shell.nav.groups.observability", "ACCESS CONTROL": "shell.nav.groups.accessControl", "DEVELOPER TOOLS": "shell.nav.groups.developerTools", @@ -553,6 +578,7 @@ export const groupText = (groupLabel: string): string => const SECTION_DISPLAY: Record = { "AI GATEWAY": "shell.nav.sections.aiGateway", + "CUSTOMER PORTAL": "shell.nav.sections.customerPortal", OBSERVABILITY: "shell.nav.sections.observability", "ACCESS CONTROL": "shell.nav.sections.accessControl", "DEVELOPER TOOLS": "shell.nav.sections.developerTools", @@ -575,7 +601,7 @@ export const labelText = (item: MenuItem): string => export const getBreadcrumb = (pathname: string): { section: string | null; title: string } => { const route = routeForPathname(pathname); if (route === "aria") return { section: null, title: "Aria Agent" }; - for (const group of menuGroups) { + for (const group of [...menuGroups, ...customerPortalGroups]) { for (const item of group.items) { const section = sectionText(group.groupLabel); if (routeOf(item) === route) return { section, title: labelText(item) }; @@ -597,6 +623,7 @@ const Sidebar_: React.FC = ({ allowVectorStoresForTeamAdmins, }) => { const { userId, accessToken, userRole, isViewOnly } = useAuthorized(); + const { mode: pluginMode } = usePluginMode(); const { t } = useTranslation(); const term = useTerminology(); const labelOf = (item: MenuItem) => (typeof item.label === "string" ? term(labelText(item)) : item.label); @@ -639,8 +666,9 @@ const Sidebar_: React.FC = ({ disableVectorStoresForInternalUsers, allowVectorStoresForTeamAdmins, }; - return visibleMenuGroups(menuGroups, context); + return visibleMenuGroups(pluginMode === "customer-portal" ? customerPortalGroups : menuGroups, context); }, [ + pluginMode, allowAgentsForTeamAdmins, allowVectorStoresForTeamAdmins, disableAgentsForInternalUsers, diff --git a/ui/litellm-dashboard/src/components/navbar.test.tsx b/ui/litellm-dashboard/src/components/navbar.test.tsx index cc90e0b8e99..38b539c51b9 100644 --- a/ui/litellm-dashboard/src/components/navbar.test.tsx +++ b/ui/litellm-dashboard/src/components/navbar.test.tsx @@ -4,6 +4,8 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, screen, waitFor, within } from "../../tests/test-utils"; import Navbar from "./navbar"; +vi.mock("next/navigation", () => ({ usePathname: () => "/api-keys", useRouter: () => ({ push: vi.fn() }) })); + // Mock the hooks and utilities vi.mock("@/components/networking", () => ({ getProxyBaseUrl: vi.fn(() => "http://localhost:4000"), diff --git a/ui/litellm-dashboard/src/contexts/PluginModeContext.tsx b/ui/litellm-dashboard/src/contexts/PluginModeContext.tsx index 380a9479829..247cae44235 100644 --- a/ui/litellm-dashboard/src/contexts/PluginModeContext.tsx +++ b/ui/litellm-dashboard/src/contexts/PluginModeContext.tsx @@ -4,7 +4,8 @@ import React, { createContext, useContext, useState, useEffect } from "react"; import { createApiClient } from "@/lib/http/client"; import { getProxyBaseUrl } from "@/components/networking"; -export type PluginMode = "ai-gateway" | string; // "ai-gateway" or a registered plugin name +export type PluginMode = "ai-gateway" | "customer-portal" | string; // built-in modes or a registered plugin name +export const BUILT_IN_MODES: ReadonlySet = new Set(["ai-gateway", "customer-portal"]); export interface PluginNavItem { key: string; @@ -84,7 +85,8 @@ export function PluginModeProvider({ children, accessToken }: PluginModeProvider // Once plugins have loaded, fall back to ai-gateway if the persisted mode is // no longer registered — including when the list came back empty (all plugins // removed). Derived rather than setState-in-effect to avoid cascading renders. - const effectiveMode = mode !== "ai-gateway" && loaded && !plugins.some((p) => p.name === mode) ? "ai-gateway" : mode; + const effectiveMode = + !BUILT_IN_MODES.has(mode) && loaded && !plugins.some((p) => p.name === mode) ? "ai-gateway" : mode; const setMode = (m: PluginMode) => { setModeState(m); diff --git a/ui/litellm-dashboard/src/i18n/locales/en.json b/ui/litellm-dashboard/src/i18n/locales/en.json index d4ba9c14998..113c77e1916 100644 --- a/ui/litellm-dashboard/src/i18n/locales/en.json +++ b/ui/litellm-dashboard/src/i18n/locales/en.json @@ -7994,14 +7994,16 @@ "observability": "OBSERVABILITY", "accessControl": "ACCESS CONTROL", "developerTools": "DEVELOPER TOOLS", - "settings": "SETTINGS" + "settings": "SETTINGS", + "customerPortal": "CUSTOMER PORTAL" }, "sections": { "aiGateway": "AI Gateway", "observability": "Observability", "accessControl": "Access Control", "developerTools": "Developer Tools", - "settings": "Settings" + "settings": "Settings", + "customerPortal": "Customer Portal" }, "items": { "home": "Home", @@ -8047,7 +8049,9 @@ "adminSettings": "Admin Settings", "costTracking": "Cost Tracking", "uiTheme": "UI Theme", - "telemetry": "Telemetry" + "telemetry": "Telemetry", + "customers": "Customers", + "chargebacks": "Chargebacks" }, "homeLink": "LiteLLM home" }, @@ -20493,5 +20497,113 @@ "removeGroup": "Remove weights for {{group}}", "summary": "Deployment Weights:", "loadFailed": "Could not load this model group's deployments. Saved weights are kept." + }, + "customerPortal": { + "switcher": { + "label": "Customer Portal", + "description": "End customers, budgets and chargebacks" + }, + "status": { + "active": "Active", + "overBudget": "Over budget", + "blocked": "Blocked" + }, + "common": { + "noBudget": "No budget", + "cancel": "Cancel", + "save": "Save", + "block": "Block", + "blockHint": "Block and Unblock change the customer record. The proxy only rejects a blocked customer's requests when the blocked_user_check callback is enabled in litellm_settings.callbacks.", + "unblock": "Unblock", + "edit": "Edit", + "total": "Total", + "loading": "Loading..." + }, + "toasts": { + "blocked": "Customer blocked", + "unblocked": "Customer unblocked", + "created": "Created {{name}}", + "updated": "Customer updated", + "createFailed": "Failed to create customer: {{error}}", + "updateFailed": "Failed to update customer: {{error}}" + }, + "list": { + "title": "Customers", + "subtitle": "End users of your product, identified by customer id on every request.", + "create": "Create Customer", + "searchPlaceholder": "Search by customer id or alias", + "columns": { + "customer": "Customer", + "spend": "Spend this month", + "budget": "Budget", + "keys": "Keys", + "status": "Status" + }, + "loading": "Loading customers...", + "empty": "No customers yet. Create one to start tracking spend.", + "budgetUsage": "{{spend}} / {{max}}" + }, + "detail": { + "spendThisMonth": "Spend this month", + "ofBudget": "of {{max}} budget", + "noBudget": "no budget", + "tpm": "{{val}} TPM", + "rpm": "{{val}} RPM", + "keysTitle": "Keys", + "keysSubtitle": "Keys that sent requests as this customer this month.", + "keysEmpty": "No traffic from this customer this month.", + "columns": { + "key": "Key", + "requests": "Requests", + "tokens": "Tokens", + "spend": "Spend" + }, + "unnamedKey": "Unnamed key", + "back": "Back to customers" + }, + "createDialog": { + "title": "Create customer", + "description": "A customer is an end user of your product. Pass its id as `user` on every request to track spend and apply its limits.", + "customerId": "Customer id", + "customerIdPlaceholder": "e.g. acme-corp", + "alias": "Alias", + "aliasPlaceholder": "Optional display name", + "maxBudget": "Monthly budget (USD)", + "maxBudgetPlaceholder": "Leave empty for no budget", + "tpm": "TPM limit", + "rpm": "RPM limit", + "noLimit": "No limit", + "limitsHelp": "Budget and limits apply to this customer across every key it uses.", + "submit": "Create" + }, + "editDialog": { + "title": "Edit customer", + "alias": "Alias", + "maxBudget": "Max budget (USD)", + "maxBudgetPlaceholder": "Leave empty for no budget", + "sharedBudgetWarning": "This customer is on budget {{budgetId}}. Changing the budget here changes it for every customer on that budget." + }, + "chargebacks": { + "title": "Chargebacks", + "subtitle": "Spend per customer for a month, from the daily spend rollup.", + "groupBy": { + "customer": "Group by customer", + "key": "Customer and key", + "model": "Customer and model" + }, + "dimension": { + "key": "Key", + "model": "Model" + }, + "exportCsv": "Export CSV", + "columns": { + "customer": "Customer", + "requests": "Requests", + "inputTokens": "Input tokens", + "outputTokens": "Output tokens", + "spend": "Spend" + }, + "empty": "No customer traffic in {{month}}." + } } } diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b229461e981..e1b0104e089 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -4549,7 +4549,6 @@ export interface paths { * - allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] - Require all user requests to use models in this specific region. * - default_model: Optional[str] - If no equivalent model in the allowed region, default all requests to this model. * - models: Optional[list[str]] - Restrict this customer's access to the listed models. - * - metadata: Optional[dict] = Metadata for customer, store information for customer. Example metadata = {"data_training_opt_out": True} * - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). * - tpm_limit: Optional[int] - [Not Implemented Yet] Specify tpm limit for a given customer (Tokens per minute) * - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute) @@ -4580,7 +4579,7 @@ export interface paths { * ``` * curl --location 'http://0.0.0.0:4000/customer/new' --header "Authorization: Bearer $LITELLM_MASTER_KEY" --header 'Content-Type: application/json' --data '{ * "user_id" : "ishaan-jaff-3", - * "allowed_region": "eu", + * "allowed_model_region": "eu", * "budget_id": "free_tier", * "models": ["gpt-4o-mini"], * "default_model": "azure/gpt-3.5-turbo-eu" @@ -5229,7 +5228,6 @@ export interface paths { * - allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] - Require all user requests to use models in this specific region. * - default_model: Optional[str] - If no equivalent model in the allowed region, default all requests to this model. * - models: Optional[list[str]] - Restrict this customer's access to the listed models. - * - metadata: Optional[dict] = Metadata for customer, store information for customer. Example metadata = {"data_training_opt_out": True} * - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). * - tpm_limit: Optional[int] - [Not Implemented Yet] Specify tpm limit for a given customer (Tokens per minute) * - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute) @@ -5260,7 +5258,7 @@ export interface paths { * ``` * curl --location 'http://0.0.0.0:4000/customer/new' --header "Authorization: Bearer $LITELLM_MASTER_KEY" --header 'Content-Type: application/json' --data '{ * "user_id" : "ishaan-jaff-3", - * "allowed_region": "eu", + * "allowed_model_region": "eu", * "budget_id": "free_tier", * "models": ["gpt-4o-mini"], * "default_model": "azure/gpt-3.5-turbo-eu"