diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 7368de1e968..d70f947469a 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -932,7 +932,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode) def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, object]: - if role == "user" or role == "system" or role == "tool": + if role in ("user", "system", "developer", "tool"): return {"type": "input_text", "text": content} else: return {"type": "output_text", "text": content} diff --git a/litellm/constants.py b/litellm/constants.py index be13d9aac5f..f5acadc32ab 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -7,6 +7,7 @@ from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_in_ran DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm")) AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview")) +AZURE_OPENAI_AUDIO_PROVIDERS: Final = frozenset({"azure", "azure_ai"}) ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5)) ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000 RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset( @@ -39,6 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset( "router_general_settings", "ignore_invalid_deployments", "fallback_access_check", + "heuristic_v2_router_limit", } ) DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index 4d043701f40..87f007ca1d5 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -18,9 +18,11 @@ caller's identity metadata, minus two things that must never be forwarded as-is: from __future__ import annotations from collections.abc import Mapping +from types import MappingProxyType from typing import Final from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, NON_INFERENCE_CALL_TYPES +from litellm.litellm_core_utils.initialize_dynamic_callback_params import initialize_standard_callback_dynamic_params from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, InternalCallOrigin BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) @@ -142,6 +144,19 @@ def forwarded_internal_call_metadata( } +def parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> Mapping[str, str]: + kwargs: Final = request_kwargs or MappingProxyType({}) + return MappingProxyType( + {k: v for k in ("litellm_session_id", "litellm_trace_id") if isinstance(v := kwargs.get(k), str)} + ) + + +def effective_turn_off_message_logging(request_kwargs: Mapping[str, object] | None) -> bool | None: + return initialize_standard_callback_dynamic_params(dict(request_kwargs) if request_kwargs else None).get( + "turn_off_message_logging" + ) + + def sanitized_forwardable_call_metadata( parent_metadata: Mapping[str, object], call_origin: InternalCallOrigin, diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 5504756ceb8..bf99035a6b1 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -5,6 +5,8 @@ Helper utilities for tracking the cost of built-in tools. from collections.abc import Mapping from typing import Final, Literal +from pydantic import ValidationError + import litellm from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS from litellm.litellm_core_utils.llm_cost_calc.utils import ( @@ -13,6 +15,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( from litellm.types.llms.openai import ( FileSearchTool, ResponsesAPIResponse, + ResponsesToolUsage, WebSearchOptions, ) from litellm.types.utils import ( @@ -32,6 +35,17 @@ def _output_item_type(output_item: object) -> str | None: return item_type if isinstance(item_type, str) else None +def _reported_web_search_requests(response_object: ResponsesAPIResponse) -> int | None: + tool_usage: Final = getattr(response_object, "tool_usage", None) + if tool_usage is None: + return None + try: + web_search: Final = ResponsesToolUsage.model_validate(tool_usage).web_search + except ValidationError: + return None + return None if web_search is None else web_search.num_requests + + def _usage_reports_server_side_web_search_calls(usage: Usage) -> bool: details: Final = getattr(usage, "server_side_tool_usage_details", None) if not isinstance(details, Mapping): @@ -182,15 +196,19 @@ class StandardBuiltInToolCostTracking: Providers that report a request count in usage (gemini, anthropic, xai, vertex) are handled by get_cost_for_web_search_request and never reach here. This path prices per call, so it must count - the web_search_call items. Chat-completions responses only expose url_citation annotations with no - count, so they floor to a single billable search. + the web_search_call items, unless the response reports the billable count itself + (Bedrock's tool_usage.web_search.num_requests, which excludes open_page fetches). Chat-completions + responses only expose url_citation annotations with no count, so they floor to a single billable search. """ - if isinstance(response_object, ResponsesAPIResponse): - count = sum( - 1 for output_item in response_object.output if _output_item_type(output_item) == "web_search_call" - ) - return max(count, 1) - return 1 + if not isinstance(response_object, ResponsesAPIResponse): + return 1 + reported: Final = _reported_web_search_requests(response_object) + if reported is not None: + return reported + count: Final = sum( + 1 for output_item in response_object.output if _output_item_type(output_item) == "web_search_call" + ) + return max(count, 1) @staticmethod def _handle_file_search_cost( diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index b34c416cd40..21587af73aa 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -428,7 +428,7 @@ def _coerce_off_peak_rate(value: object, default: float) -> float: return default -def _apply_off_peak_pricing( +def apply_off_peak_pricing( model_info: ModelInfo, current_time: datetime | None, prompt_base_cost: float, @@ -462,7 +462,7 @@ def _apply_off_peak_to_base_costs( has no field for them. """ prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs - off_peak_prompt, off_peak_completion, off_peak_cache_read = _apply_off_peak_pricing( + off_peak_prompt, off_peak_completion, off_peak_cache_read = apply_off_peak_pricing( model_info, current_time, prompt, completion, cache_read ) return (off_peak_prompt, off_peak_completion, cache_creation, cache_creation_above_1hr, off_peak_cache_read) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 355faa41e68..6ae17bac6ff 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2337,6 +2337,9 @@ class CustomStreamWrapper: else: self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() + if self.stream_options is None: + usage: Final = calculate_total_usage(chunks=self.chunks) + processed_chunk._hidden_params["usage"] = usage # pyright: ignore[reportPrivateUsage] # sync parity # see sync __next__'s sibling branch: deliberately do NOT restore # here - this chunk is still this call's own data, and restoring # before returning it would corrupt the caller's own log diff --git a/litellm/llms/azure/chat/gpt_5_transformation.py b/litellm/llms/azure/chat/gpt_5_transformation.py index 6fdd277a04f..3189f5b57ac 100644 --- a/litellm/llms/azure/chat/gpt_5_transformation.py +++ b/litellm/llms/azure/chat/gpt_5_transformation.py @@ -7,6 +7,7 @@ from litellm.exceptions import UnsupportedParamsError from litellm.llms.openai.chat.gpt_5_transformation import ( OpenAIGPT5Config, _get_effort_level, + is_gpt_reasoning_series_name, ) from litellm.types.llms.openai import AllMessageValues @@ -35,26 +36,7 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): @classmethod def is_model_gpt_5_model(cls, model: str) -> bool: - """Check if the Azure model string refers to a gpt-5 variant. - - Accepts both explicit gpt-5 model names and the ``gpt5_series/`` prefix - used for manual routing. - """ - # The gpt-5-chat* family (gpt-5-chat, gpt-5-chat-latest, gpt-5-chat-2025-08-07, - # …) are regular chat models: they support temperature and tool_choice but NOT - # reasoning_effort. They must NOT be routed through the GPT-5 reasoning path. - # - # Versioned chat models such as gpt-5.3-chat and gpt-5.1-chat ARE reasoning - # models and must stay on the GPT-5 path. The distinguishing feature is that - # the gpt-5-chat family has a literal "-chat" immediately after "gpt-5" - # (i.e. "gpt-5-chat…"), while versioned chat models interpose a minor version - # number (i.e. "gpt-5.-chat"). - # - # Using a startswith("gpt-5-chat") prefix check on the normalized name (rather - # than a substring check) makes this boundary explicit and avoids any ambiguity - # if future model names coincidentally contain "gpt-5-chat" as an interior run. - _normalized: Final = model.split("/")[-1] # strip provider prefix, e.g. "azure/" - return ("gpt-5" in model and not _normalized.startswith("gpt-5-chat")) or "gpt5_series" in model + return is_gpt_reasoning_series_name(model) or "gpt5_series" in model def get_supported_openai_params(self, model: str) -> list[str]: """Get supported parameters for Azure OpenAI GPT-5 models. diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 0ac0662205a..880a51eb584 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_azure_openai_messages, ) from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.openai.chat.gpt_5_transformation import GPT_REASONING_SERIES_MARKERS from litellm.types.llms.azure import ( API_VERSION_MONTH_SUPPORTED_RESPONSE_FORMAT, API_VERSION_YEAR_SUPPORTED_RESPONSE_FORMAT, @@ -139,7 +140,7 @@ class AzureOpenAIConfig(BaseConfig): name family needs the rename, including the ``gpt-5-chat*`` models that are excluded from the reasoning path by https://github.com/BerriAI/litellm/issues/13781. """ - return "gpt-5" in model or "gpt5_series" in model + return any(marker in model for marker in GPT_REASONING_SERIES_MARKERS) or "gpt5_series" in model def _is_response_format_supported_model(self, model: str) -> bool: """ diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 7fe9d3dec52..f2d405e9a17 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, ) from litellm.llms.azure.common_utils import BaseAzureLLM +from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error from litellm.llms.openai.openai import OpenAIConfig @@ -207,20 +208,18 @@ class AzureAIStudioConfig(OpenAIConfig): message["content"] = texts return stripped_messages - def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool: - try: - if "/" in model: - model = model.split("/", 1)[1] - if ( - model in litellm.open_ai_chat_completion_models - or model in litellm.open_ai_text_completion_models - or model in litellm.open_ai_embedding_models - ): - return True + def _is_foundry_model_inference_base(self, api_base: str) -> bool: + return is_foundry_model_inference_base(api_base) - except Exception: + def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool: + if api_base is None or self._is_foundry_model_inference_base(api_base): return False - return False + stripped_model: Final = model.split("/", 1)[1] if "/" in model else model + return ( + stripped_model in litellm.open_ai_chat_completion_models + or stripped_model in litellm.open_ai_text_completion_models + or stripped_model in litellm.open_ai_embedding_models + ) def _get_openai_compatible_provider_info( self, diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 26a90157455..aa34bab5b2e 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -1,5 +1,6 @@ from collections.abc import Mapping from typing import Final, Literal +from urllib.parse import urlparse import litellm from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter @@ -10,6 +11,14 @@ from litellm.types.router import GenericLiteLLMParams AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"] +def is_foundry_model_inference_base(api_base: str) -> bool: + parsed: Final = urlparse(api_base) + host: Final = parsed.hostname + if host is None or not host.endswith(".services.ai.azure.com"): + return False + return "/openai/deployments" not in parsed.path + + def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None: """ Resolve an Entra ID / OAuth access token for an Azure AI Foundry deployment. diff --git a/litellm/llms/azure_ai/embed/handler.py b/litellm/llms/azure_ai/embed/handler.py index 65c3997c099..c65edbf56e6 100644 --- a/litellm/llms/azure_ai/embed/handler.py +++ b/litellm/llms/azure_ai/embed/handler.py @@ -1,8 +1,10 @@ from typing import Final +from urllib.parse import urlsplit, urlunsplit from openai import OpenAI import litellm +from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -16,6 +18,16 @@ from litellm.utils import convert_to_model_response_object from .cohere_transformation import AzureAICohereConfig +def _foundry_models_route_base(api_base: str | None) -> str | None: + if api_base is None or not is_foundry_model_inference_base(api_base): + return api_base + parts: Final = urlsplit(api_base) + path: Final = parts.path.rstrip("/") + if path.endswith("/models"): + return api_base + return urlunsplit((parts.scheme, parts.netloc, f"{path}/models", parts.query, parts.fragment)) + + class AzureAIEmbedding(OpenAIChatCompletion): def _process_response( self, @@ -214,6 +226,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): assemble result in-order, and return """ + resolved_api_base: Final = _foundry_models_route_base(api_base) if aembedding is True: return self.async_embedding( model, @@ -223,7 +236,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): model_response, optional_params, api_key, - api_base, + resolved_api_base, client, ) @@ -245,7 +258,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): model_response=model_response, optional_params=optional_params, api_key=api_key, - api_base=api_base, + api_base=resolved_api_base, client=client, ) @@ -262,7 +275,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): model_response, optional_params, api_key, - api_base, + resolved_api_base, client=(client if client is not None and isinstance(client, OpenAI) else None), aembedding=aembedding, shared_session=shared_session, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 71a598a6fe7..2896d889f96 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -99,9 +99,12 @@ from litellm.types.containers.main import ( ) from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig from litellm.types.integrations.custom_logger import ( + NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, AgenticLoopPlan, AgenticLoopRequestPatch, AgenticLoopSafetyError, + converted_stream_requested, + is_interception_internal_key, ) from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, @@ -2760,6 +2763,7 @@ class BaseLLMHTTPHandler: ) if self._has_agentic_completion_hook(logging_obj): + agentic_kwargs: Final = dict(litellm_params) # mutable-ok: agentic hooks mutate kwargs in place final_response: Final = run_async_function( self._call_agentic_completion_hooks, response=initial_response, @@ -2770,10 +2774,19 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, stream=False, custom_llm_provider=custom_llm_provider, - kwargs=dict(litellm_params), + kwargs=agentic_kwargs, api_surface="responses", ) - return final_response if final_response is not None else initial_response + result: Final = final_response if final_response is not None else initial_response + if converted_stream_requested(agentic_kwargs) and not agentic_kwargs.get("_agentic_loop_depth"): + return self._wrap_responses_response_as_fake_stream( + result=result, + model=model, + responses_api_provider_config=responses_api_provider_config, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + ) + return result return initial_response @@ -2939,6 +2952,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, ) + agentic_kwargs: Final = dict(litellm_params) # mutable-ok: agentic hooks mutate kwargs in place final_response: Final = await self._call_agentic_completion_hooks( response=initial_response, model=model, @@ -2948,15 +2962,12 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, stream=False, custom_llm_provider=custom_llm_provider, - kwargs=dict(litellm_params), + kwargs=agentic_kwargs, api_surface="responses", ) result: Final = final_response if final_response is not None else initial_response - interception_converted_stream: Final = litellm_params.get( - "_code_interpreter_interception_converted_stream" - ) or litellm_params.get("_websearch_interception_converted_stream") - if interception_converted_stream and not litellm_params.get("_agentic_loop_depth"): + if converted_stream_requested(agentic_kwargs) and not agentic_kwargs.get("_agentic_loop_depth"): return self._wrap_responses_response_as_fake_stream( result=result, model=model, @@ -5420,8 +5431,7 @@ class BaseLLMHTTPHandler: kwargs_for_followup: Final = { k: v for k, v in kwargs.items() - if not k.startswith("_websearch_interception") - and not k.startswith("_compression_interception") + if not is_interception_internal_key(k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES) and k != "_code_interpreter_interception_converted_stream" and k not in internal_keys and k not in optional_params diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index dd5bee1fe8b..d8eb1f9f8d7 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -7,11 +7,13 @@ cached, cache-creation, output, reasoning) is billed at that one tier's rate. See https://help.aliyun.com/zh/model-studio/billing-for-model-studio """ -from dataclasses import dataclass +from dataclasses import dataclass, replace +from datetime import datetime from typing import Final from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate from litellm.litellm_core_utils.llm_cost_calc.utils import ( + apply_off_peak_pricing, parse_completion_tokens_details, parse_prompt_tokens_details, ) @@ -32,6 +34,19 @@ class TokenBreakdown: return self.text_tokens + self.cached_tokens + self.cache_creation_tokens +@dataclass(frozen=True, slots=True) +class TokenRates: + input_rate: float + cache_read_rate: float + cache_creation_rate: float + output_rate: float + reasoning_rate: float | None + + @property + def billed_reasoning_rate(self) -> float: + return self.output_rate if self.reasoning_rate is None else self.reasoning_rate + + def _extract_token_breakdown(usage: Usage) -> TokenBreakdown: prompt_details: Final = parse_prompt_tokens_details(usage) cached_tokens: Final = prompt_details["cache_hit_tokens"] @@ -57,69 +72,75 @@ def _flat_rate(model_info: ModelInfo, cost_key: str, fallback_cost_key: str) -> return float(value) -def _calculate_prompt_cost( - breakdown: TokenBreakdown, - model_info: ModelInfo, - tier: dict | None, -) -> float: - if tier is not None: - return ( - (breakdown.text_tokens * tier_rate(tier, "input_cost_per_token")) - + (breakdown.cached_tokens * tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token")) - + ( - breakdown.cache_creation_tokens - * tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token") - ) - ) - - input_cost: Final = float(model_info.get("input_cost_per_token") or 0.0) - cache_read_cost: Final = _flat_rate(model_info, "cache_read_input_token_cost", "input_cost_per_token") - cache_creation_cost: Final = _flat_rate(model_info, "cache_creation_input_token_cost", "input_cost_per_token") - - return ( - (breakdown.text_tokens * input_cost) - + (breakdown.cached_tokens * cache_read_cost) - + (breakdown.cache_creation_tokens * cache_creation_cost) +def _flat_rates(model_info: ModelInfo) -> TokenRates: + reasoning_rate: Final = model_info.get("output_cost_per_reasoning_token") + return TokenRates( + input_rate=float(model_info.get("input_cost_per_token") or 0.0), + cache_read_rate=_flat_rate(model_info, "cache_read_input_token_cost", "input_cost_per_token"), + cache_creation_rate=_flat_rate(model_info, "cache_creation_input_token_cost", "input_cost_per_token"), + output_rate=float(model_info.get("output_cost_per_token") or 0.0), + reasoning_rate=None if reasoning_rate is None else float(reasoning_rate), ) -def _calculate_completion_cost( - breakdown: TokenBreakdown, - model_info: ModelInfo, - tier: dict | None, -) -> float: +def _tier_rates(model_info: ModelInfo, tier: dict) -> TokenRates: # A tier that declares output rates keeps the request on them, all-or-nothing. A tier table # spelling out only input rates would serve every completion for free, so there the model's # own output rates stand in - tier_declares_output: Final = tier is not None and "output_cost_per_token" in tier - output_cost: Final = ( - tier_rate(tier, "output_cost_per_token") - if tier_declares_output - else float(model_info.get("output_cost_per_token") or 0.0) - ) - tier_declares_reasoning: Final = tier is not None and "output_cost_per_reasoning_token" in tier - model_reasoning_rate: Final = None if tier_declares_output else model_info.get("output_cost_per_reasoning_token") - reasoning_cost: Final = ( - tier_rate(tier, "output_cost_per_reasoning_token", "output_cost_per_token") - if tier_declares_reasoning - else float(model_reasoning_rate) - if model_reasoning_rate is not None - else output_cost + flat_rates: Final = _flat_rates(model_info) + tier_declares_output: Final = "output_cost_per_token" in tier + tier_declares_reasoning: Final = "output_cost_per_reasoning_token" in tier + return TokenRates( + input_rate=tier_rate(tier, "input_cost_per_token"), + cache_read_rate=tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token"), + cache_creation_rate=tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token"), + output_rate=tier_rate(tier, "output_cost_per_token") if tier_declares_output else flat_rates.output_rate, + reasoning_rate=( + tier_rate(tier, "output_cost_per_reasoning_token") + if tier_declares_reasoning + else None + if tier_declares_output + else flat_rates.reasoning_rate + ), ) - return (breakdown.completion_tokens * output_cost) + (breakdown.reasoning_tokens * reasoning_cost) + +def _off_peak_rates(model_info: ModelInfo, current_time: datetime | None, rates: TokenRates) -> TokenRates: + input_rate, output_rate, cache_read_rate = apply_off_peak_pricing( + model_info, current_time, rates.input_rate, rates.output_rate, rates.cache_read_rate + ) + return replace(rates, input_rate=input_rate, output_rate=output_rate, cache_read_rate=cache_read_rate) -def cost_per_token(model: str, usage: Usage, custom_llm_provider: str = "dashscope") -> tuple[float, float]: +def _bill(breakdown: TokenBreakdown, rates: TokenRates) -> tuple[float, float]: + prompt_cost: Final = ( + (breakdown.text_tokens * rates.input_rate) + + (breakdown.cached_tokens * rates.cache_read_rate) + + (breakdown.cache_creation_tokens * rates.cache_creation_rate) + ) + completion_cost: Final = (breakdown.completion_tokens * rates.output_rate) + ( + breakdown.reasoning_tokens * rates.billed_reasoning_rate + ) + return prompt_cost, completion_cost + + +def cost_per_token( + model: str, + usage: Usage, + custom_llm_provider: str = "dashscope", + current_time: datetime | None = None, +) -> tuple[float, float]: """ Calculate cost per token for Dashscope models. - Supports both tiered and flat pricing with cached and reasoning tokens. + Supports both tiered and flat pricing with cached and reasoning tokens, and swaps in the + model's off_peak_pricing rates while one of its windows is open. Args: model: Model name without provider prefix usage: LiteLLM Usage block custom_llm_provider: The provider id the request resolved to; dashscope or one of its brand aliases + current_time: The moment the request is billed at; defaults to now, UTC Returns: Tuple[float, float] - (prompt_cost_in_usd, completion_cost_in_usd) @@ -133,8 +154,7 @@ def cost_per_token(model: str, usage: Usage, custom_llm_provider: str = "dashsco if tiered_pricing else None ) + standard_rates: Final = _flat_rates(model_info) if tier is None else _tier_rates(model_info, tier) + rates: Final = _off_peak_rates(model_info, current_time, standard_rates) - prompt_cost: Final = _calculate_prompt_cost(breakdown=breakdown, model_info=model_info, tier=tier) - completion_cost: Final = _calculate_completion_cost(breakdown=breakdown, model_info=model_info, tier=tier) - - return prompt_cost, completion_cost + return _bill(breakdown, rates) diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index 0223be300b0..b02f953425d 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -61,6 +61,14 @@ def _get_effort_level(value: str | dict | None) -> str | None: return None +GPT_REASONING_SERIES_MARKERS: Final = ("gpt-5", "gpt-6") + + +def is_gpt_reasoning_series_name(model: str) -> bool: + normalized: Final = model.split("/")[-1] + return any(marker in model for marker in GPT_REASONING_SERIES_MARKERS) and not normalized.startswith("gpt-5-chat") + + class OpenAIGPT5Config(OpenAIGPTConfig): """Configuration for gpt-5 models including GPT-5-Codex variants. @@ -73,21 +81,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig): @classmethod def is_model_gpt_5_model(cls, model: str) -> bool: - # The gpt-5-chat* family (gpt-5-chat, gpt-5-chat-latest, gpt-5-chat-2025-08-07, - # …) are regular chat models: they support temperature and tool_choice but NOT - # reasoning_effort. They must NOT be routed through the GPT-5 reasoning path. - # - # Versioned chat models such as gpt-5.3-chat and gpt-5.1-chat ARE reasoning - # models and must stay on the GPT-5 path. The distinguishing feature is that - # the gpt-5-chat family has a literal "-chat" immediately after "gpt-5" - # (i.e. "gpt-5-chat…"), while versioned chat models interpose a minor version - # number (i.e. "gpt-5.-chat"). - # - # Using a startswith("gpt-5-chat") prefix check on the normalized name (rather - # than a substring check) makes this boundary explicit and avoids any ambiguity - # if future model names coincidentally contain "gpt-5-chat" as an interior run. - _normalized: Final = model.split("/")[-1] # strip provider prefix, e.g. "openai/" - return "gpt-5" in model and not _normalized.startswith("gpt-5-chat") + return is_gpt_reasoning_series_name(model) @classmethod def is_model_gpt_5_search_model(cls, model: str) -> bool: @@ -122,6 +116,8 @@ class OpenAIGPT5Config(OpenAIGPTConfig): def is_model_gpt_5_4_plus_model(cls, model: str) -> bool: """Check if the model is gpt-5.4 or newer (5.4, 5.5, 5.6, etc., including pro).""" model_name: Final = model.split("/")[-1] + if model_name.startswith("gpt-6"): + return True if not model_name.startswith("gpt-5."): return False try: diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index bcd4ea43243..2db6d78a218 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -11,6 +11,7 @@ import time import uuid from collections.abc import AsyncIterator, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Optional +from urllib.parse import urlsplit import httpx import openai @@ -43,6 +44,14 @@ _OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(OpenAI) _AZURE_OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(AzureOpenAI) +_OPENAI_API_HOST: Final[str] = "api.openai.com" + + +def is_openai_backed_api_base(api_base: str) -> bool: + hostname: Final = urlsplit(api_base).hostname + return hostname is not None and (hostname == _OPENAI_API_HOST or hostname.endswith(f".{_OPENAI_API_HOST}")) + + class OpenAIError(BaseLLMException): def __init__( self, diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 1cfc6e06ee9..edc8d64d9c2 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -2,7 +2,6 @@ import time import types from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast -from urllib.parse import urlparse import httpx @@ -55,6 +54,7 @@ from .common_utils import ( OpenAIError, build_output_token_limit_response, drop_params_from_unprocessable_entity_error, + is_openai_backed_api_base, is_output_token_limit_error, ) from .workload_identity import resolve_openai_workload_identity_config @@ -1190,10 +1190,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): """ if stream_options is not None: return {"stream_options": stream_options} - else: - # by default litellm will include usage for openai endpoints - if api_base is None or urlparse(api_base).hostname == "api.openai.com": - return {"stream_options": {"include_usage": True}} + if api_base is None or is_openai_backed_api_base(api_base): + return {"stream_options": {"include_usage": True}} return {} # Embedding diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index a0db7aadb9e..ecc89c5f135 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -33,8 +33,9 @@ import time import uuid from collections.abc import Mapping, Sequence from dataclasses import dataclass +from itertools import accumulate from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Union, cast +from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from pydantic import BaseModel, TypeAdapter @@ -42,6 +43,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, OpenAiResponsesToChatCompletionStreamIterator, ) from litellm.llms.base_llm.guardrail_translation.base_translation import ( @@ -74,6 +76,7 @@ from litellm.types.llms.openai import ( OutputTextDoneEvent, ResponseAPIUsage, ResponseCompletedEvent, + ResponsesAPIOptionalRequestParams, ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, @@ -115,6 +118,199 @@ class ResponsesStreamChunk(TypedDict, total=False): content_index: ReadOnly[int] +_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + {"function_call_output": "output", "message": "content"} +) + +_EMPTY_RESPONSES_REQUEST: Final[ResponsesAPIOptionalRequestParams] = {} + + +def _item_rewrite_field(item: Mapping[str, object]) -> str | None: + item_type: Final = item.get("type") + if item_type is None: + return "content" if "content" in item else None + if not isinstance(item_type, str): + return None + return _PATCHABLE_ITEM_FIELDS.get(item_type) + + +def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapping[str, object] | None: + field: Final = _item_rewrite_field(item) + if field is None or not isinstance(rewritten, Mapping): + return None + rewritten_content: Final = rewritten.get("content") + if isinstance(item.get(field), str) and isinstance(rewritten_content, str): + return {**item, field: rewritten_content} # mutable-ok: request input items must stay JSON-plain dicts + rewritten_row: Final = cast("AllMessageValues", rewritten) # cast-ok: guardrails hand back chat-shaped rows + converted_items, _ = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api( + [rewritten_row] # mutable-ok: converter signature takes a list + ) + if len(converted_items) != 1 or not isinstance(converted_items[0], Mapping): + return None + first_converted: Final = cast("Mapping[str, object]", converted_items[0]) # cast-ok: isinstance-checked above + converted_value: Final = first_converted.get(field) + if converted_value is None: + return None + return {**item, field: converted_value} # mutable-ok: request input items must stay JSON-plain dicts + + +def _is_function_call_item(item: object) -> bool: + return isinstance(item, Mapping) and item.get("type") in ("function_call", "custom_tool_call") + + +def _last_message_role(messages: Sequence[object]) -> str | None: + if not messages: + return None + last: Final = messages[-1] + role: Final = last.get("role") if isinstance(last, Mapping) else getattr(last, "role", None) + return role if isinstance(role, str) else None + + +def _provenance_unit_bounds( + raw_input: Sequence[object], + solo_conversions: Sequence[Sequence[object]], +) -> tuple[tuple[int, int], ...]: + trailing_roles: Final = tuple( + accumulate( + (_last_message_role(messages) for messages in solo_conversions), + lambda previous, current: current if current is not None else previous, + ) + ) + start_indexes: Final = tuple( + index + for index in range(len(raw_input)) + if index == 0 or not (_is_function_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant") + ) + return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input)))) + + +def _input_item_provenance( + raw_input: Sequence[object], + expected_messages: Sequence[object], +) -> tuple[Mapping[int, int], frozenset[int]] | None: + if not all(isinstance(item, Mapping) for item in raw_input): + return None + solo_conversions: Final = tuple( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast("ResponseInputParam", [item]), # cast-ok: items checked as Mappings above + responses_api_request=_EMPTY_RESPONSES_REQUEST, + ) + for item in raw_input + ) + full_conversion: Final = tuple( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast("ResponseInputParam", list(raw_input)), # cast-ok: items checked as Mappings above + responses_api_request=_EMPTY_RESPONSES_REQUEST, + ) + ) + if full_conversion != tuple(expected_messages): + return None + units: Final = _provenance_unit_bounds(raw_input, solo_conversions) + unit_messages: Final = tuple( + tuple(solo_conversions[start]) + if end - start == 1 + else tuple( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast("ResponseInputParam", list(raw_input[start:end])), # cast-ok: checked as Mappings above + responses_api_request=_EMPTY_RESPONSES_REQUEST, + ) + ) + for start, end in units + ) + if tuple(message for messages in unit_messages for message in messages) != full_conversion: + return None + boundaries: Final = tuple(accumulate((len(messages) for messages in unit_messages), initial=0)) + item_for_message: Final = MappingProxyType( + { + message_index: start + for unit_index, (start, end) in enumerate(units) + if end - start == 1 + for message_index in range(boundaries[unit_index], boundaries[unit_index + 1]) + } + ) + tainted: Final = frozenset( + message_index + for unit_index, (start, end) in enumerate(units) + if end - start > 1 + for message_index in range(boundaries[unit_index], boundaries[unit_index + 1]) + ) + return item_for_message, tainted + + +class _RequestFields(NamedTuple): + input: tuple[object, ...] + instructions: str | None + + +class _ExtractedInputs(NamedTuple): + inputs: GenericGuardrailAPIInputs + task_mappings: tuple[tuple[int, int | None], ...] + + +def _patched_request_fields( + raw_input: object, + instructions: object, + original_messages: Sequence[object], + structured_messages: Sequence[object], +) -> _RequestFields | None: + if not isinstance(raw_input, list) or len(original_messages) != len(structured_messages): + return None + offset: Final = 1 if instructions else 0 + provenance: Final = _input_item_provenance(raw_input, tuple(original_messages)[offset:]) + if provenance is None: + return None + item_for_message, tainted = provenance + changed: Final = tuple( + (index, rewritten) + for index, (original, rewritten) in enumerate(zip(original_messages, structured_messages)) + if original != rewritten + ) + instruction_rewrites: Final = tuple(rewritten for index, rewritten in changed if index < offset) + rewritten_instructions: Final = ( + instruction_rewrites[0].get("content") + if instruction_rewrites and isinstance(instruction_rewrites[0], Mapping) + else instructions + ) + instructions_value: Final = rewritten_instructions if isinstance(rewritten_instructions, str) else None + if rewritten_instructions is not None and instructions_value is None: + return None + body_changes: Final = tuple((index - offset, rewritten) for index, rewritten in changed if index >= offset) + if any(message_index in tainted or message_index not in item_for_message for message_index, _ in body_changes): + return None + replacements: Final = MappingProxyType( + { + item_for_message[message_index]: _rewritten_input_item( + cast("Mapping[str, object]", raw_input[item_for_message[message_index]]), # cast-ok: checked Mappings + rewritten, + ) + for message_index, rewritten in body_changes + } + ) + if len(replacements) != len(body_changes) or any(item is None for item in replacements.values()): + return None + return _RequestFields( + input=tuple(replacements.get(index, item) for index, item in enumerate(raw_input)), + instructions=instructions_value, + ) + + +def _patch_or_convert_request_fields( + raw_input: object, + instructions: object, + original_messages: Sequence[object], + structured_messages: Sequence[AllMessageValues], +) -> _RequestFields | None: + if not isinstance(structured_messages, list): + return None + patched: Final = _patched_request_fields(raw_input, instructions, original_messages, structured_messages) + if patched is not None: + return patched + input_items, converted_instructions = ( + LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(structured_messages) + ) + return _RequestFields(input=tuple(input_items), instructions=converted_instructions) + + def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int: sequence_numbers: Final = ( item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None) @@ -162,9 +358,8 @@ class OpenAIResponsesHandler(BaseTranslation): Handles both string input and list of message objects. """ input_data: Final[str | ResponseInputParam | None] = data.get("input") - if input_data is None: + if not isinstance(input_data, (str, list)): return data - structured_messages: Final = self.get_structured_messages(data) raw_tools: Final = data.get("tools") original_tools: Final[tuple[Mapping[str, object], ...]] = ( @@ -173,94 +368,93 @@ class OpenAIResponsesHandler(BaseTranslation): flattened_tool_groups: Final = tuple( form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original_tools) ) - flattened_tools: Final = tuple( - cast(ChatCompletionToolParam, tool) # cast-ok: mcp tools ride along in the guardrail's tool list - for group in flattened_tool_groups - for tool in group - ) - tools_to_check: Final[list[ChatCompletionToolParam]] = list( # mutable-ok: guardrail inputs want a list - copy.deepcopy(flattened_tools) - ) - - # Handle simple string input - if isinstance(input_data, str): - inputs = GenericGuardrailAPIInputs(texts=[input_data]) - if tools_to_check: - inputs["tools"] = tools_to_check - if structured_messages: - inputs["structured_messages"] = structured_messages - # Include model information if available - model = data.get("model") - if model: - inputs["model"] = model - - guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs=inputs, - request_data=data, - input_type="request", - logging_obj=litellm_logging_obj, - ) - guardrailed_texts = guardrailed_inputs.get("texts", []) - data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data - self._apply_guardrailed_tools_to_data( - data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") - ) - verbose_proxy_logger.debug("OpenAI Responses API: Processed string input") - return data - - # Handle list input (ResponseInputParam) - if not isinstance(input_data, list): + extracted: Final = self._extract_guardrail_inputs(data, input_data, flattened_tool_groups) + if not extracted.inputs.get("texts"): return data + if structured_messages: + extracted.inputs["structured_messages"] = structured_messages + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( + inputs=extracted.inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + self._apply_guardrailed_tools_to_data( + data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") + ) + written_back: Final = self._written_back_request_fields(data, structured_messages, guardrailed_inputs) + if written_back is not None: + data["input"] = list(written_back.input) # mutable-ok: JSON body + if written_back.instructions is None: + data.pop("instructions", None) + else: + data["instructions"] = written_back.instructions # rebind-ok: data is an out-param + elif isinstance(input_data, str): + guardrailed_texts: Final = guardrailed_inputs.get("texts") or () + data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data # rebind-ok: data is an out-param + else: + await self._apply_guardrail_responses_to_input( + messages=input_data, + responses=guardrailed_inputs.get("texts") or (), + task_mappings=extracted.task_mappings, + ) + verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", data.get("input")) + return data + def _extract_guardrail_inputs( + self, + data: Mapping[str, object], + input_data: "str | ResponseInputParam", + flattened_tool_groups: Sequence[Sequence[Mapping[str, object]]], + ) -> _ExtractedInputs: texts_to_check: Final[list[str]] = [] images_to_check: Final[list[str]] = [] task_mappings: Final[list[tuple[int, int | None]]] = [] - - # Step 1: Extract all text content, images, and tools - for msg_idx, message in enumerate(input_data): - self._extract_input_text_and_images( - message=message, - msg_idx=msg_idx, - texts_to_check=texts_to_check, - images_to_check=images_to_check, - task_mappings=task_mappings, + tools_to_check: Final[list[ChatCompletionToolParam]] = list( # mutable-ok: guardrail inputs want a list + copy.deepcopy( + tuple( + cast(ChatCompletionToolParam, tool) # cast-ok: mcp tools ride along in the guardrail's tool list + for group in flattened_tool_groups + for tool in group + ) ) + ) + if isinstance(input_data, str): + texts_to_check.append(input_data) + else: + for msg_idx, message in enumerate(input_data): + self._extract_input_text_and_images( + message=message, + msg_idx=msg_idx, + texts_to_check=texts_to_check, + images_to_check=images_to_check, + task_mappings=task_mappings, + ) + inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) + if images_to_check: + inputs["images"] = images_to_check + if tools_to_check: + inputs["tools"] = tools_to_check + model: Final = data.get("model") + if isinstance(model, str): + inputs["model"] = model + return _ExtractedInputs(inputs=inputs, task_mappings=tuple(task_mappings)) - # Step 2: Apply guardrail to all texts in batch - if texts_to_check: - inputs = GenericGuardrailAPIInputs(texts=texts_to_check) - if images_to_check: - inputs["images"] = images_to_check - if tools_to_check: - inputs["tools"] = tools_to_check - if structured_messages: - inputs["structured_messages"] = structured_messages - # Include model information if available - model = data.get("model") - if model: - inputs["model"] = model - guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs=inputs, - request_data=data, - input_type="request", - logging_obj=litellm_logging_obj, - ) - - guardrailed_texts = guardrailed_inputs.get("texts", []) - self._apply_guardrailed_tools_to_data( - data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") - ) - - # Step 3: Map guardrail responses back to original input structure - await self._apply_guardrail_responses_to_input( - messages=input_data, - responses=guardrailed_texts, - task_mappings=task_mappings, - ) - - verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", input_data) - - return data + @staticmethod + def _written_back_request_fields( + data: Mapping[str, object], + structured_messages: Sequence[AllMessageValues] | None, + guardrailed_inputs: GenericGuardrailAPIInputs, + ) -> _RequestFields | None: + guardrailed: Final = guardrailed_inputs.get("structured_messages") + if guardrailed is None or guardrailed is structured_messages: + return None + return _patch_or_convert_request_fields( + data.get("input"), + data.get("instructions"), + structured_messages or (), + guardrailed, + ) def extract_request_tool_names(self, data: dict) -> list[str]: """Extract tool names from Responses API request (tools[].name for function @@ -331,8 +525,8 @@ class OpenAIResponsesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, messages: Any, # Can be List[Dict[str, Any]] or ResponseInputParam - responses: list[str], - task_mappings: list[tuple[int, int | None]], + responses: Sequence[str], + task_mappings: Sequence[tuple[int, int | None]], ) -> None: """ Apply guardrail responses back to input messages. diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 01313e95878..b97521b90c2 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo ) from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name from litellm.responses.litellm_completion_transformation.custom_tools import TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import * @@ -88,7 +89,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): parts: Final = model.split("/") if len(parts) > 1 and parts[0] not in ("openai",): return False - return "gpt-5" in model and "gpt-5-chat" not in model + return is_gpt_reasoning_series_name(model) @staticmethod def _supports_reasoning_effort_none(model: str) -> bool: diff --git a/litellm/main.py b/litellm/main.py index 0bca4a7350e..2929790f2bd 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -26,6 +26,7 @@ from copy import deepcopy from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args +from urllib.parse import urlsplit from litellm._logging import _redact_string from litellm._uuid import uuid @@ -60,6 +61,7 @@ if TYPE_CHECKING: from litellm.types.utils import TokenCountResponse from litellm.constants import ( + AZURE_OPENAI_AUDIO_PROVIDERS, DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, ) @@ -984,6 +986,12 @@ def mock_completion( _OPENAI_DEFAULT_API_BASE: Final = "https://api.openai.com/v1" +_OPENAI_API_HOST: Final = "api.openai.com" + + +def _is_openai_backed_api_base(api_base: str) -> bool: + hostname: Final = urlsplit(api_base).hostname + return hostname is not None and (hostname == _OPENAI_API_HOST or hostname.endswith(f".{_OPENAI_API_HOST}")) def _resolve_openai_api_base(api_base: str | None) -> str: @@ -1053,7 +1061,7 @@ def responses_api_bridge_check( # 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 (the default OpenAI endpoint, or Azure OpenAI where api_base is + # 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 @@ -1068,14 +1076,15 @@ def responses_api_bridge_check( reasoning_active = reasoning_effort.get("effort") != "none" or reasoning_effort.get("summary") is not None else: reasoning_active = reasoning_effort != "none" - # The reasoning+tools constraint is enforced only by the real OpenAI endpoint (and Azure OpenAI). - # Resolve the effective base arg>global>env>default exactly as the chat handler does, so a custom - # base set via litellm.api_base or OPENAI_BASE_URL/OPENAI_API_BASE isn't misread as the default and - # bridged to a /responses route it lacks. A whitespace-only base collapses to the default too. - resolved_api_base: Final = _resolve_openai_api_base(api_base) - on_constraint_enforcing_endpoint: Final = custom_llm_provider == "azure" or resolved_api_base.strip() in ( - "", - _OPENAI_DEFAULT_API_BASE, + # The reasoning+tools constraint is enforced by the real OpenAI backend behind any api.openai.com + # host (the default URL or a PrivateLink hostname such as .privatelink.api.openai.com) and + # by Azure OpenAI. Resolve the effective base arg>global>env>default exactly as the chat handler + # does, so a custom base set via litellm.api_base or OPENAI_BASE_URL/OPENAI_API_BASE isn't misread + # as the default and bridged to a /responses route it lacks. A whitespace-only base collapses to + # the default too. + resolved_api_base: Final = _resolve_openai_api_base(api_base).strip() + on_constraint_enforcing_endpoint: Final = ( + custom_llm_provider == "azure" or resolved_api_base == "" or _is_openai_backed_api_base(resolved_api_base) ) if ( custom_llm_provider in ("openai", "azure") @@ -7769,7 +7778,7 @@ def transcription( provider=LlmProviders(custom_llm_provider), ) - if custom_llm_provider == "azure" and provider_config is None: + if custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS and provider_config is None: # azure configs api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") @@ -8056,7 +8065,10 @@ def speech( custom_llm_provider=custom_llm_provider, ) response: HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent] | None = None - if custom_llm_provider == "openai" or custom_llm_provider in litellm.openai_compatible_providers: + if custom_llm_provider == "openai" or ( + custom_llm_provider in litellm.openai_compatible_providers + and custom_llm_provider not in AZURE_OPENAI_AUDIO_PROVIDERS + ): if voice is None or not (isinstance(voice, str)): raise litellm.BadRequestError( message="'voice' is required to be passed as a string for OpenAI TTS", @@ -8110,7 +8122,7 @@ def speech( aspeech=aspeech, shared_session=shared_session, ) - elif custom_llm_provider == "azure": + elif custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS: # Check if this is Azure Speech Service (Cognitive Services TTS) if model.startswith("speech/"): from litellm.llms.azure.text_to_speech.transformation import ( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3b6e7912253..1c51385318b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29354,6 +29354,75 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, + "gpt-6-astra": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, + "cache_creation_input_token_cost_flex": 6.25e-06, + "cache_creation_input_token_cost_priority": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 4e-06, + "cache_read_input_token_cost_flex": 5e-07, + "cache_read_input_token_cost_priority": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "input_cost_per_token_above_272k_tokens_flex": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 4e-05, + "input_cost_per_token_batches": 5e-06, + "input_cost_per_token_flex": 5e-06, + "input_cost_per_token_priority": 2e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, + "output_cost_per_token_above_272k_tokens_priority": 0.00015, + "output_cost_per_token_batches": 2.5e-05, + "output_cost_per_token_flex": 2.5e-05, + "output_cost_per_token_priority": 0.0001, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "gpt-5.6": { "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, @@ -53485,6 +53554,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53519,6 +53593,11 @@ "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, "output_cost_per_token": 1.32e-05, "output_cost_per_token_above_272k_tokens": 1.98e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53581,6 +53660,11 @@ "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, "output_cost_per_token": 1.32e-06, "output_cost_per_token_above_272k_tokens": 1.98e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53769,6 +53853,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53800,6 +53889,11 @@ "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "output_cost_per_token": 1.65e-05, "output_cost_per_token_above_272k_tokens": 2.475e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0aea72be1e2..5d5a25e7cd6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1216,7 +1216,7 @@ class GenerateKeyRequest(KeyRequestBase): organization_id: str | None = None project_id: str | None = None - @field_validator("team_id", "organization_id", mode="before") + @field_validator("team_id", "organization_id", "project_id", mode="before") @classmethod def treat_cleared_id_as_unset(cls, v: object) -> object: if v == "": @@ -4239,6 +4239,8 @@ class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase): access_group_id: str access_group_name: str models: tuple[str, ...] + mcp_server_ids: tuple[str, ...] = () + agent_ids: tuple[str, ...] = () class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py index 677f1a0fdda..55bb1e3925a 100644 --- a/litellm/proxy/auth/litellm_license.py +++ b/litellm/proxy/auth/litellm_license.py @@ -16,6 +16,10 @@ if TYPE_CHECKING: from litellm.proxy._types import EnterpriseLicenseData +AUTO_ROUTER_LICENSE_FEATURE: Final = "auto_router" +HEURISTIC_V2_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit." + + class LicenseCheck: """ - Check if license in env @@ -149,6 +153,19 @@ class LicenseCheck: return False return team_count > _max_teams_in_license + def heuristic_v2_router_limit(self) -> int | None: + """ + How many heuristic_v2 auto-routers this proxy may hold: unlimited (None) only when the + signed license lists the auto_router feature, otherwise one. A license verified through + the API carries no feature list, so it does not lift the limit either. + """ + if self.airgapped_license_data is None: + return 1 + allowed_features: Final = self.airgapped_license_data.get("allowed_features") + if isinstance(allowed_features, list) and AUTO_ROUTER_LICENSE_FEATURE in allowed_features: + return None + return 1 + def verify_license_without_api_request(self, public_key, license_key): try: from cryptography.hazmat.primitives import hashes @@ -179,19 +196,21 @@ class LicenseCheck: # Decode and parse the data license_data: Final = json.loads(message.decode()) - self.airgapped_license_data = EnterpriseLicenseData(**license_data) - # debug information provided in license data verbose_proxy_logger.debug("License data: %s", license_data) # Check expiration date expiration_date: Final = datetime.strptime(license_data["expiration_date"], "%Y-%m-%d") if expiration_date < datetime.now(): + self.airgapped_license_data = None return False, "License has expired" + self.airgapped_license_data = EnterpriseLicenseData(**license_data) + return True except Exception as e: + self.airgapped_license_data = None verbose_proxy_logger.debug( "litellm.proxy.auth.litellm_license.py::verify_license_without_api_request - Unable to verify License locally. - %s", e, diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py index 1fc3c06e6bf..93d859066b0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -916,9 +916,10 @@ class CompresrGuardrail(CustomGuardrail): def _mirror_texts_channel(input_texts: object, applied: _CompressionResult) -> list[object] | None: """Compressed content mirrored into the Responses `texts` channel. - The chat/Anthropic handlers round-trip ``structured_messages``; the - Responses translation cannot rebuild its input from chat messages and - instead writes back through ``texts``. This matches by value, so a + The chat/Anthropic/Responses handlers round-trip + ``structured_messages``; translations without that round-trip write + back through ``texts``, so the compressed content is mirrored there + too. This matches by value, so a replacement is applied only when it is unambiguous: one compression per text, and every occurrence in ``texts`` accounted for by a compressed target. Anything else is left uncompressed rather than risk a wrong or diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index fc881a60f43..9d993384461 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -50,6 +50,9 @@ if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel BYPASS_HEADER: Final = "x-headroom-bypass" +_STREAM_CONVERTIBLE_CALL_TYPES: Final = frozenset( + (CallTypes.completion, CallTypes.acompletion, CallTypes.responses, CallTypes.aresponses) +) HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve" _HASH_PATTERN: Final = re.compile(r"hash=([a-f0-9]{24})") _HASH_CACHE_TTL_SECONDS: Final = 15 * 60 @@ -725,6 +728,10 @@ class HeadroomGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Headroom: %s header set; skipping compression", BYPASS_HEADER) return inputs + if request_data.get("background"): + verbose_proxy_logger.debug("Headroom: background request; skipping compression") + return inputs + structured_messages: Final = inputs.get("structured_messages") if not _is_object_list(structured_messages) or not structured_messages: return inputs @@ -826,9 +833,9 @@ class HeadroomGuardrail(CustomGuardrail): ) -> dict[str, Any] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict base_result: Final = await super().async_pre_call_deployment_hook(kwargs, call_type) effective: Final = base_result if base_result is not None else kwargs - if call_type not in (CallTypes.completion, CallTypes.acompletion): + if call_type not in _STREAM_CONVERTIBLE_CALL_TYPES: return base_result - if not effective.get("stream"): + if not effective.get("stream") or effective.get("background"): return base_result if not has_headroom_retrieve_tool(effective.get("tools")): return base_result diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 613e726f89d..82ee33cbc39 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -13,10 +13,11 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import datetime import json -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence +from contextlib import AbstractAsyncContextManager, asynccontextmanager from json import JSONDecodeError from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeVar, cast from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator @@ -49,6 +50,7 @@ from litellm.proxy._types import ( TeamModelDeleteRequest, UserAPIKeyAuth, ) +from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, @@ -96,6 +98,9 @@ from litellm.router_strategy.complexity_router import ( from litellm.router_utils.auto_router_model_naming import ( STRATEGY_ROUTER_PARAM_FIELDS, carries_complexity_router_settings, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, + uses_heuristic_v2_classifier, validate_complexity_router_config_placement, validate_complexity_router_config_write, validate_strategy_router_model_write, @@ -153,6 +158,8 @@ class _ProxyModelTable(Protocol): def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ... + def create(self, *, data: Mapping[str, object]) -> Awaitable[_ProxyModelRow]: ... + def update( self, *, where: Mapping[str, object], data: Mapping[str, object] ) -> Awaitable[_ProxyModelRow | None]: ... @@ -166,6 +173,9 @@ class _TxModelTables(Protocol): litellm_proxymodeltable: _ProxyModelTable +_RowT = TypeVar("_RowT") + + class _ExistingModelRow(Protocol): @property def litellm_params(self) -> Mapping[str, object]: ... @@ -269,6 +279,66 @@ def _raise_on_strategy_router_write_violation( ) +HEURISTIC_V2_SLOT_LOCK_KEY: Final = 5_872_301 +_HEURISTIC_V2_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" +_HEURISTIC_V2_DB_ROWS_SQL: Final = """ +SELECT count(*)::int AS held FROM "LiteLLM_ProxyModelTable" +WHERE model_id <> $1 + AND (CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END) + -> 'complexity_router_config' ->> 'classifier_type' = 'heuristic_v2' +""" + + +def _effective_complexity_router_config( + incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None +) -> object: + """The complexity config a write leaves on the row: the incoming one when the write carries it, else the stored one.""" + incoming: Final = None if incoming_params is None else incoming_params.complexity_router_config + if incoming is not None or existing_params is None: + return incoming + return existing_params.complexity_router_config + + +@asynccontextmanager +async def _heuristic_v2_slot( + prisma_client: PrismaClient, *, effective_config: object, model_id: str | None +) -> AsyncGenerator[_ProxyModelTable, None]: + """Hand out the model table to write through while the row's claim on a heuristic_v2 slot is settled. + + A write that leaves the row on classifier_type heuristic_v2 under a limited license runs + inside one transaction that takes an advisory lock in its own statement before counting + (a statement's snapshot predates anything it locks), so pods cannot both pass the count: + the DB rows (any pod, either JSON shape) plus this proxy's config.yaml routers are judged + against the license limit and the write is refused with a 403 before it happens. The row + being edited keeps its own slot through ``model_id``. Every other write, and every write on + an unlimited license, goes through the repository table with no lock. Only the row write + itself may run inside: anything that needs a second connection (the team model bookkeeping) + must wait until the transaction has committed and the lock is released. The transaction + writes bypass the repository's publish-on-write, so the config change is published once + after commit, the way delete_team_models does. + """ + from litellm.proxy.proxy_server import _license_check, llm_router + + limit: Final = _license_check.heuristic_v2_router_limit() + if limit is None or not uses_heuristic_v2_classifier(effective_config): + yield _proxy_model_table(prisma_client) + return + async with prisma_client.db.tx() as tx_ctx: + tables: Final[_TxModelTables] = tx_ctx + await tx_ctx.query_raw(_HEURISTIC_V2_LOCK_SQL, HEURISTIC_V2_SLOT_LOCK_KEY) + rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw(_HEURISTIC_V2_DB_ROWS_SQL, model_id or "") + db_held: Final = rows[0].get("held") if rows else 0 + config_rows: Final = () if llm_router is None else tuple(llm_router.config_deployments()) + held: Final = (db_held if isinstance(db_held, int) else 0) + count_heuristic_v2_routers(config_rows) + violation: Final = heuristic_v2_limit_violation(held=held + 1, limit=limit) + if violation is not None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {HEURISTIC_V2_LICENSE_REMEDY}" + ) + yield tables.litellm_proxymodeltable + await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") + + ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add" _REQUIRED_RATE_LIMIT_FIELDS: Final = ("rpm", "tpm") @@ -720,22 +790,29 @@ async def patch_model( ) requested_model_name: Final = patch_data.model_name + stored_model_name: str | None = None + + async def write_row(update_data: PrismaCompatibleUpdateDBModel) -> _ProxyModelRow | None: + nonlocal stored_model_name + stored_model_name = update_data.get("model_name") + update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name + update_data["updated_at"] = cast(str, get_utc_datetime()) + async with _heuristic_v2_slot( + prisma_client, + effective_config=_effective_complexity_router_config( + patch_data.litellm_params, db_model.litellm_params + ), + model_id=model_id, + ) as table: + return await table.update(where={"model_id": model_id}, data=update_data) + # Handle team model updates with proper alias management - update_data: Final = await _update_team_model_in_db( + updated_model: Final = await _update_team_model_in_db( db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, - ) - - # Add metadata about update - update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name - update_data["updated_at"] = cast(str, get_utc_datetime()) - - # Perform partial update - updated_model: Final = await _proxy_model_table(prisma_client).update( - where={"model_id": model_id}, - data=update_data, + write_row=write_row, ) if updated_model is None: @@ -746,7 +823,6 @@ async def patch_model( param=None, ) - stored_model_name: Final = update_data.get("model_name") if ( stored_model_name is not None and stored_model_name == requested_model_name @@ -980,7 +1056,8 @@ async def _add_model_to_db( prisma_client: PrismaClient, new_encryption_key: str | None = None, should_create_model_in_db: bool = True, -) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None": + slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None, +) -> "_ProxyModelRow | LiteLLM_ProxyModelTable": # encrypt litellm params # _litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True) _original_litellm_model_name: Final = model_params.litellm_params.model @@ -998,18 +1075,20 @@ async def _add_model_to_db( if model_params.model_info.id is not None: _data["model_id"] = model_params.model_info.id _create_data: Final = cast("Mapping[str, object]", _data) # cast-ok: str-keyed json payload built just above - if should_create_model_in_db: - model_response = await ModelRepository(prisma_client).table.create(data=_create_data) - else: - model_response = LiteLLM_ProxyModelTable(**_data) - return model_response + if not should_create_model_in_db: + return LiteLLM_ProxyModelTable(**_data) + if slot is None: + return await _proxy_model_table(prisma_client).create(data=_create_data) + async with slot as table: + return await table.create(data=_create_data) async def _add_team_model_to_db( model_params: Deployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, -) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None": + slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None, +) -> "_ProxyModelRow | LiteLLM_ProxyModelTable": """ If 'team_id' is provided, @@ -1040,6 +1119,7 @@ async def _add_team_model_to_db( model_params=model_params, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, + slot=slot, ) if original_model_name: @@ -1060,7 +1140,8 @@ async def _update_team_model_in_db( patch_data: updateDeployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, -) -> PrismaCompatibleUpdateDBModel: + write_row: Callable[[PrismaCompatibleUpdateDBModel], Awaitable[_RowT]], +) -> _RowT: """ Handle team model updates with proper alias management. @@ -1068,6 +1149,9 @@ async def _update_team_model_in_db( - Creates unique internal model_name and team alias - Adds model to team object - Preserves team_public_model_name for external reference + + The row is written through ``write_row`` before the team's model list is touched, so a + refused or failed write leaves the team as it was (the create path orders itself the same way). """ # Validate team_id if present in patch_data from litellm.proxy.proxy_server import premium_user @@ -1079,9 +1163,7 @@ async def _update_team_model_in_db( premium_user=premium_user, ) - # Validated before any write, beside the premium check the create path already runs - # here. The team ACL is updated below and autocommits, so a validator that raises - # further down would leave the team mutated and the deployment row never written. + # Validated before the row write, beside the premium check the create path already runs here. # # The merged view is what gets stored, so that is what has to satisfy the invariants. # Validating the patch alone rejected a partial edit of an already valid deployment: @@ -1101,7 +1183,7 @@ async def _update_team_model_in_db( # No team_id in patch, proceed with standard update if patch_team_id is None: - return update_db_model(db_model=db_model, updated_patch=patch_data) + return await write_row(update_db_model(db_model=db_model, updated_patch=patch_data)) # Determine public model name public_model_name: Final = _get_public_model_name( @@ -1120,11 +1202,14 @@ async def _update_team_model_in_db( db_team_id: Final = db_model.model_info.team_id if db_model.model_info else None is_new_team_assignment: Final = db_team_id != patch_team_id + # Team rows keep their internal UUID-based model_name; the public name lives in model_info + patch_data.model_name = f"model_name_{patch_team_id}_{uuid.uuid4()}" if is_new_team_assignment else None + row: Final = await write_row(update_db_model(db_model=db_model, updated_patch=patch_data)) + if is_new_team_assignment: await _setup_new_team_model_assignment( team_id=patch_team_id, public_model_name=public_model_name, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, ) else: @@ -1132,12 +1217,11 @@ async def _update_team_model_in_db( team_id=patch_team_id, public_model_name=public_model_name, db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, ) - return update_db_model(db_model=db_model, updated_patch=patch_data) + return row def _get_public_model_name( @@ -1189,13 +1273,9 @@ def _get_public_model_name( async def _setup_new_team_model_assignment( team_id: str, public_model_name: str, - patch_data: updateDeployment, user_api_key_dict: UserAPIKeyAuth, ) -> None: - """Set up a new team model with unique name and team membership.""" - unique_model_name: Final = f"model_name_{team_id}_{uuid.uuid4()}" - patch_data.model_name = unique_model_name - + """Register a newly team-assigned model's public name on the team.""" await team_model_add( data=TeamModelAddRequest( team_id=team_id, @@ -1385,7 +1465,6 @@ async def _update_existing_team_model_assignment( team_id: str, public_model_name: str, db_model: Deployment, - patch_data: updateDeployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient | None, ) -> None: @@ -1409,9 +1488,6 @@ async def _update_existing_team_model_assignment( old_public_name: Final = db_model.model_info.team_public_model_name if db_model.model_info else None if old_public_name and public_model_name != old_public_name: - # Clear user-supplied public name from patch before any early return so the - # caller does not overwrite the internal UUID-based model_name in the DB. - patch_data.model_name = None if prisma_client is None: verbose_proxy_logger.warning( "prisma_client not initialized; skipping public name update entirely to avoid orphaned entries" @@ -1459,10 +1535,6 @@ async def _update_existing_team_model_assignment( # else: old_public_name == public_model_name (no rename needed) # No team_model_add/delete calls required; public name is already registered - # Always clear patch_data.model_name to prevent caller from overwriting - # the internal UUID-based model_name in the DB with the user-supplied public name - patch_data.model_name = None - class ModelManagementAuthChecks: """ @@ -1878,18 +1950,19 @@ async def add_new_model( reload_outcome: ReconcileOutcome = ReconcileOutcome(still_desired=None, live_after=None) try: _original_litellm_model_name: Final = model_params.model_name - if model_params.model_info.team_id is None: - model_response = await _add_model_to_db( - model_params=priced_model_params, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - ) - else: - model_response = await _add_team_model_to_db( - model_params=priced_model_params, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - ) + add_model: Final = ( + _add_model_to_db if model_params.model_info.team_id is None else _add_team_model_to_db + ) + model_response = await add_model( + model_params=priced_model_params, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + slot=_heuristic_v2_slot( + prisma_client, + effective_config=priced_model_params.litellm_params.complexity_router_config, + model_id=priced_model_params.model_info.id, + ), + ) reload_outcome = await proxy_config.add_deployment( prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj ) @@ -1903,6 +1976,8 @@ async def add_new_model( passed_model_info=priced_model_params.model_info, ) except Exception as e: + if isinstance(e, HTTPException): + raise verbose_proxy_logger.exception("Exception in add_new_model: %s", e) else: @@ -2070,10 +2145,17 @@ async def update_model( "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, **({} if renamed_to is None else {"model_name": renamed_to}), } - model_response: Final = await _proxy_model_table(prisma_client).update( - where={"model_id": _model_id}, - data=_data, - ) + async with _heuristic_v2_slot( + prisma_client, + effective_config=_effective_complexity_router_config( + model_params.litellm_params, deployment.litellm_params + ), + model_id=_model_id, + ) as table: + model_response: Final = await table.update( + where={"model_id": _model_id}, + data=_data, + ) if renamed_to is not None: await sync_access_groups_for_renamed_model( prisma_client=prisma_client, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 714cf252e69..72cdc75c29e 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -4318,6 +4318,8 @@ async def _resolve_team_access_group_resources( access_group_id=group.access_group_id, access_group_name=group.access_group_name, models=tuple(group.access_model_names or ()), + mcp_server_ids=tuple(group.access_mcp_server_ids or ()), + agent_ids=tuple(group.access_agent_ids or ()), ) for group in resolved_groups ), diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index b48b8d81494..29f216fd450 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1730,6 +1730,16 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict: return headers +def _is_vertex_anthropic_count_tokens_route(endpoint: str) -> bool: + return endpoint.rsplit("/", 1)[-1].split(":", 1)[0] == "count-tokens" + + +def _upstream_headers_for_vertex_route(endpoint: str, headers: Mapping[str, str]) -> Mapping[str, str]: + if not _is_vertex_anthropic_count_tokens_route(endpoint): + return headers + return MappingProxyType({name: value for name, value in headers.items() if name.lower() != "anthropic-beta"}) + + def get_vertex_pass_through_handler( call_type: Literal["discovery", "aiplatform"], # noqa: UP037 # ruff reports quoted Literal values here ) -> BaseVertexAIPassThroughHandler: @@ -2128,7 +2138,7 @@ async def _base_vertex_proxy_route( endpoint_func: Final = create_pass_through_route( endpoint=endpoint, target=target, - custom_headers=headers, + custom_headers=_upstream_headers_for_vertex_route(endpoint, headers), is_streaming_request=is_streaming_request, ) # dynamically construct pass-through endpoint based on incoming path diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 27132c90e05..96b425f2a24 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -120,6 +120,8 @@ from litellm.router_utils.add_retry_fallback_headers import ( from litellm.router_utils.auto_router_model_naming import ( STRATEGY_ROUTER_PARAM_FIELDS, carries_complexity_router_settings, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, validate_complexity_router_config_placement, ) from litellm.types.utils import ( @@ -301,7 +303,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.fallback_model_access import router_fallback_access_check from litellm.proxy.auth.handle_jwt import JWTHandler -from litellm.proxy.auth.litellm_license import LicenseCheck +from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY, LicenseCheck from litellm.proxy.auth.model_checks import ( expand_wildcard_deployments_for_model_info, get_all_fallbacks, @@ -4316,6 +4318,19 @@ def validate_deployment_complexity_router_placement(model: Mapping[str, object]) raise ValueError(f"model {model.get('model_name', '')!r}: {violation}") +def validate_heuristic_v2_router_limit(model_list: Sequence[Mapping[str, object]], *, limit: int | None) -> None: + """ + Refuse to start when config.yaml defines more heuristic_v2 auto-routers than the license allows. + + Checked here rather than left to router registration for the same reason as the two + validators above: the proxy builds its router with `ignore_invalid_deployments=True`, so + the router's own refusal would turn the extra router into a silently missing model. + """ + violation: Final = heuristic_v2_limit_violation(held=count_heuristic_v2_routers(model_list), limit=limit) + if violation is not None: + raise ValueError(f"config.yaml model_list: {violation} {HEURISTIC_V2_LICENSE_REMEDY}") + + def pin_complexity_router_model_id(model: dict) -> None: # mutable-ok: out-param, model_info is stamped in place """ Stamps `model_info.id` from the raw litellm_params before plugin resolution swaps @@ -5721,6 +5736,7 @@ class ProxyConfig: model_list: Final = config.get("model_list", None) if model_list: router_params["model_list"] = model_list + validate_heuristic_v2_router_limit(model_list, limit=_license_check.heuristic_v2_router_limit()) print( # noqa: T201 "\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m" ) @@ -5810,6 +5826,7 @@ class ProxyConfig: ), ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid fallback_access_check=router_fallback_access_check, + heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit, ) if redis_usage_cache is not None and router.cache.redis_cache is None: @@ -6270,6 +6287,7 @@ class ProxyConfig: search_tools=search_tools, ignore_invalid_deployments=True, fallback_access_check=router_fallback_access_check, + heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit, ) verbose_proxy_logger.debug("updated llm_router: %s", llm_router) else: diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 295faaa980a..2a50d5170f0 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -173,6 +173,9 @@ class _SessionSpendRow(TypedDict): session_cache_hit_count: ReadOnly[int] session_llm_count: ReadOnly[int] session_agent_count: ReadOnly[int] + session_total_prompt_tokens: ReadOnly[int] + session_total_completion_tokens: ReadOnly[int] + session_total_tokens: ReadOnly[int] session_models: ReadOnly[Sequence[str]] @@ -188,6 +191,9 @@ class _SessionSpendStats(NamedTuple): session_cache_hit_count: int session_llm_count: int session_agent_count: int + session_total_prompt_tokens: int + session_total_completion_tokens: int + session_total_tokens: int session_models: Sequence[str] session_models_truncated: bool @@ -4287,8 +4293,8 @@ async def _build_ui_spend_logs_response( Build the paginated response for the UI spend-logs endpoint. When ``enrich_session_counts`` is ``True`` (the default for the v1/UI - endpoint), each row is enriched with ``session_total_count`` plus spend - and call-type aggregates so the frontend knows which sessions are + endpoint), each row is enriched with ``session_total_count`` plus spend, + token and call-type aggregates so the frontend knows which sessions are expandable (multi-call sessions). One ``GROUP BY (session_id, api_key)`` query serves every referenced session, keyed per api key so two callers reusing a session id never see each other's totals. Rows without a @@ -4356,7 +4362,10 @@ async def _build_ui_spend_logs_response( COUNT(*) FILTER ( WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL} )::int AS session_llm_count, - COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count + COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count, + COALESCE(SUM(prompt_tokens), 0)::bigint AS session_total_prompt_tokens, + COALESCE(SUM(completion_tokens), 0)::bigint AS session_total_completion_tokens, + COALESCE(SUM(total_tokens), 0)::bigint AS session_total_tokens FROM "LiteLLM_SpendLogs" WHERE session_id = ANY($1::text[]) AND api_key = ANY($2::text[]) @@ -4389,6 +4398,9 @@ async def _build_ui_spend_logs_response( session_cache_hit_count=int(row.get("session_cache_hit_count") or 0), session_llm_count=int(row.get("session_llm_count") or 0), session_agent_count=int(row.get("session_agent_count") or 0), + session_total_prompt_tokens=int(row.get("session_total_prompt_tokens") or 0), + session_total_completion_tokens=int(row.get("session_total_completion_tokens") or 0), + session_total_tokens=int(row.get("session_total_tokens") or 0), session_models=models[:_SESSION_MODELS_LIMIT], session_models_truncated=len(models) > _SESSION_MODELS_LIMIT, ) @@ -4418,6 +4430,9 @@ async def _build_ui_spend_logs_response( row_dict["session_cache_hit_count"] = session_stats.session_cache_hit_count row_dict["session_llm_count"] = session_stats.session_llm_count row_dict["session_agent_count"] = session_stats.session_agent_count + row_dict["session_total_prompt_tokens"] = session_stats.session_total_prompt_tokens + row_dict["session_total_completion_tokens"] = session_stats.session_total_completion_tokens + row_dict["session_total_tokens"] = session_stats.session_total_tokens row_dict["session_models"] = session_stats.session_models row_dict["session_models_truncated"] = session_stats.session_models_truncated enriched.append(row_dict) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cab2bd6d9db..24a65457928 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -7531,6 +7531,9 @@ def create_model_info_response( max_input_tokens = configured_input if configured_output is not None: max_output_tokens = configured_output + configured_mode: Final = llm_router.get_configured_mode(model_id) + if isinstance(configured_mode, str): + base["mode"] = configured_mode if max_input_tokens is not None: base["max_input_tokens"] = max_input_tokens diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index c928398a87f..fe4732c6492 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -30,7 +30,10 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user -from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store +from litellm.proxy.vector_store_endpoints.utils import ( + can_user_access_vector_store, + filter_listable_vector_stores, +) from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ManagedVectorStoresRepository from litellm.types.vector_stores import ( @@ -390,11 +393,10 @@ async def list_vector_stores( # Filter vector stores based on access control accessible_vector_stores: Final = [] - for vs in vector_store_map.values(): - if await _check_vector_store_access(vs, user_api_key_dict): - redacted = LiteLLM_ManagedVectorStore(**vs) - redacted["litellm_params"] = _redact_sensitive_litellm_params(vs.get("litellm_params")) - accessible_vector_stores.append(redacted) + for vs in await filter_listable_vector_stores(vector_store_map.values(), user_api_key_dict): + redacted = LiteLLM_ManagedVectorStore(**vs) + redacted["litellm_params"] = _redact_sensitive_litellm_params(vs.get("litellm_params")) + accessible_vector_stores.append(redacted) total_count: Final = len(accessible_vector_stores) total_pages: Final = (total_count + page_size - 1) // page_size diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 93f1510bf22..6e94a5a88ac 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -1,11 +1,17 @@ import json import re +from collections.abc import Iterable +from types import MappingProxyType from typing import Any, Final, Literal from fastapi import HTTPException, Request import litellm from litellm._logging import verbose_proxy_logger +from litellm.proxy._experimental.mcp_server.ui_session_utils import ( + is_ui_session_credential, + resolve_ui_session_team_ids, +) from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, LitellmUserRoles, @@ -160,10 +166,16 @@ async def can_user_access_vector_store( if _is_proxy_admin(user_api_key_dict): return True - vector_store_team_id: Final = vector_store.get("team_id") - if vector_store_team_id is None: + if vector_store.get("team_id") is None: return True + return await _is_vector_store_granted(vector_store, user_api_key_dict) + + +async def _is_vector_store_granted( + vector_store: LiteLLM_ManagedVectorStore, + user_api_key_dict: UserAPIKeyAuth, +) -> bool: vector_store_id: Final = vector_store.get("vector_store_id") or "" key_object_permission = user_api_key_dict.object_permission @@ -178,12 +190,70 @@ async def can_user_access_vector_store( if _object_permission_allows_vector_store(team_object_permission, vector_store_id): return True - if user_api_key_dict.team_id is not None and user_api_key_dict.team_id == vector_store_team_id: - return True + return user_api_key_dict.team_id is not None and user_api_key_dict.team_id == vector_store.get("team_id") + +async def _team_auth_context(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> UserAPIKeyAuth: + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + team: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + return user_api_key_dict.model_copy( + update=MappingProxyType( + { + "team_id": team_id, + "team_object_permission": team.object_permission, + "team_object_permission_id": team.object_permission_id, + } + ) + ) + + +async def _vector_store_listing_auth_contexts( + user_api_key_dict: UserAPIKeyAuth, +) -> tuple[UserAPIKeyAuth, ...]: + if not is_ui_session_credential(user_api_key_dict): + return (user_api_key_dict,) + session_key_context: Final = user_api_key_dict.model_copy( + update=MappingProxyType({"team_id": None, "team_object_permission": None, "team_object_permission_id": None}) + ) + team_ids: Final = await resolve_ui_session_team_ids(user_api_key_dict) + team_contexts: Final = tuple([await _team_auth_context(team_id, user_api_key_dict) for team_id in team_ids]) + return (session_key_context, *team_contexts) + + +async def _is_vector_store_granted_to_any( + vector_store: LiteLLM_ManagedVectorStore, + auth_contexts: tuple[UserAPIKeyAuth, ...], +) -> bool: + for auth_context in auth_contexts: + if await _is_vector_store_granted(vector_store, auth_context): + return True return False +async def filter_listable_vector_stores( + vector_stores: Iterable[LiteLLM_ManagedVectorStore], + user_api_key_dict: UserAPIKeyAuth, +) -> tuple[LiteLLM_ManagedVectorStore, ...]: + """Non-admins only see stores their key, one of their teams' object_permission, or team ownership grants.""" + if _is_proxy_admin(user_api_key_dict): + return tuple(vector_stores) + + auth_contexts: Final = await _vector_store_listing_auth_contexts(user_api_key_dict) + return tuple([vs for vs in vector_stores if await _is_vector_store_granted_to_any(vs, auth_contexts)]) + + async def get_litellm_managed_vector_store( vector_store_id: str, ) -> LiteLLM_ManagedVectorStore | None: diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index d4b9f4e8cce..3862aec445f 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -8,6 +8,7 @@ from typing import Any, Final, Literal, cast import litellm from litellm.constants import ( + AZURE_OPENAI_AUDIO_PROVIDERS, REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, request_timeout, @@ -400,7 +401,7 @@ async def _arealtime( litellm_metadata=_build_litellm_metadata(kwargs), query_params=query_params, ) - elif _custom_llm_provider == "azure": + elif _custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS: api_base = dynamic_api_base or litellm_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # set API KEY api_key = dynamic_api_key or litellm.api_key or litellm.openai_key or get_secret_str("AZURE_API_KEY") diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index bc25f4fffb1..d7f8cd8f8bd 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -562,6 +562,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): hidden_params: Final = getattr(chunk, "_hidden_params", None) if hidden_params is not None: chunk_dict["_hidden_params"] = dict(hidden_params) if isinstance(hidden_params, dict) else hidden_params + if ( + chunk_dict.get("usage") is None + and isinstance(hidden_params, dict) + and hidden_params.get("usage") is not None + ): + chunk_dict["usage"] = hidden_params["usage"] return chunk_dict def create_reasoning_summary_text_done_event( diff --git a/litellm/router.py b/litellm/router.py index 2b8b342d253..dc750941559 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -21,7 +21,7 @@ import time import traceback import weakref from collections import defaultdict -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Iterator, Mapping, Sequence from functools import lru_cache, partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast @@ -117,6 +117,9 @@ from litellm.router_utils.add_retry_fallback_headers import ( from litellm.router_utils.auto_router_model_naming import ( AUTO_ROUTER_MODEL_PREFIX, classify_strategy_router_model, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, + uses_heuristic_v2_classifier, ) from litellm.router_utils.batch_utils import ( _get_router_metadata_variable_name, @@ -211,6 +214,7 @@ from litellm.types.router import ( DeploymentTypedDict, FallbackAccessCheck, GuardrailTypedDict, + HeuristicV2RouterLimit, LiteLLM_Params, MockRouterTestingParams, ModelGroupInfo, @@ -683,6 +687,7 @@ class Router: background_health_check_model_groups: Sequence[str] | None = None, enable_weighted_failover: bool = False, fallback_access_check: FallbackAccessCheck | None = None, + heuristic_v2_router_limit: HeuristicV2RouterLimit | None = None, ) -> None: """ Initialize the Router class with the given parameters for caching, reliability, and routing strategy. @@ -759,6 +764,7 @@ class Router: self.set_verbose = set_verbose self.ignore_invalid_deployments = ignore_invalid_deployments + self.heuristic_v2_router_limit = heuristic_v2_router_limit self.fallback_access_check: Final = fallback_access_check self.debug_level = debug_level self.enable_pre_call_checks = enable_pre_call_checks @@ -8796,6 +8802,30 @@ class Router: """ return classify_strategy_router_model(litellm_params.model) == "complexity" + def config_deployments(self) -> Iterator[Mapping[str, object]]: + """The model_list rows that came from config.yaml rather than the DB (``model_info.db_model`` unset).""" + for deployment in self.model_list: + if not isinstance(deployment, Mapping): + continue + model_info = deployment.get("model_info") + if not (isinstance(model_info, Mapping) and model_info.get("db_model")): + yield deployment + + def heuristic_v2_router_limit_violation(self) -> str | None: + """ + Why one more heuristic_v2 router cannot join this router, or None when it can. + + Judged against every deployment currently on the model_list; an upsert pops the row being + edited first, so an edit of an existing heuristic_v2 router keeps its own slot. The limit is + resolved on every call through ``heuristic_v2_router_limit``; unset means unlimited, which + is the SDK default, and the proxy injects a resolver backed by its license. + """ + limit: Final = self.heuristic_v2_router_limit() if self.heuristic_v2_router_limit is not None else None + others: Final = count_heuristic_v2_routers( + deployment for deployment in self.model_list if isinstance(deployment, Mapping) + ) + return heuristic_v2_limit_violation(held=others + 1, limit=limit) + def init_complexity_router_deployment(self, deployment: Deployment): """ Initialize the complexity-router deployment. @@ -8813,6 +8843,10 @@ class Router: ) complexity_router_config: Final[dict | None] = deployment.litellm_params.complexity_router_config + if uses_heuristic_v2_classifier(complexity_router_config): + limit_violation: Final = self.heuristic_v2_router_limit_violation() + if limit_violation is not None: + raise ValueError(limit_violation) default_model: str | None = deployment.litellm_params.complexity_router_default_model @@ -9636,8 +9670,16 @@ class Router: raise e def _restore_deployment_after_failed_upsert(self, previous_deployment: Deployment | None, model_id: str) -> None: + """Put a deployment back the way it was before a failed upsert popped it. + + A rollback re-admits state that was already serving, so it does not go through the + heuristic_v2 ceiling a newcomer gets: with the ceiling tightened since the deployment first + registered, judging the rollback would drop a serving router over an unrelated failed edit. + """ if previous_deployment is None or self.has_model_id(model_id): return + limit_resolver: Final = self.heuristic_v2_router_limit + self.heuristic_v2_router_limit = None try: self.add_deployment(deployment=previous_deployment) verbose_router_logger.info( @@ -9652,6 +9694,8 @@ class Router: model_id, restore_error, ) + finally: + self.heuristic_v2_router_limit = limit_resolver @staticmethod def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]: @@ -9981,6 +10025,17 @@ class Router: coerce_token_limit(model_info.get("max_output_tokens")), ) + def get_configured_mode(self, model_name: str) -> "str | None": + """Return the mode explicitly configured for a concrete deployment.""" + deployment: Final = self.get_deployment_by_model_group_name(model_group_name=model_name) + if deployment is None: + return None + + mode: Final = deployment.model_info.get("mode") + if isinstance(mode, str) and mode.strip(): + return mode + return None + def get_configured_display_name(self, model_name: str) -> "str | None": """ Return the display_name explicitly configured in a concrete deployment's diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index c77745a498d..6b443026f61 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -2,23 +2,41 @@ Auto-Routing Strategy that works with a Semantic Router Config """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Optional +from pydantic import BaseModel, ConfigDict + from litellm._logging import verbose_router_logger from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.internal_call_metadata import ( + effective_turn_off_message_logging, + forwarded_internal_call_metadata, + parent_session_kwargs, +) +from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN if TYPE_CHECKING: from semantic_router.routers import SemanticRouter from semantic_router.routers.base import Route from litellm.router import Router + from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder from litellm.types.router import PreRoutingHookResponse else: Router = Any PreRoutingHookResponse = Any Route = Any SemanticRouter = Any + LiteLLMRouterEncoder = Any + + +class _CallerMetadata(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + metadata: Mapping[str, object] | None = None + litellm_metadata: Mapping[str, object] | None = None class AutoRouter(CustomLogger): @@ -50,6 +68,8 @@ class AutoRouter(CustomLogger): """ from semantic_router.routers import SemanticRouter + from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder + self.auto_router_config_path: str | None = auto_router_config_path self.auto_router_config: str | None = auto_router_config self.auto_sync_value = self.DEFAULT_AUTO_SYNC_VALUE @@ -59,6 +79,11 @@ class AutoRouter(CustomLogger): self.embedding_model: str = embedding_model self.max_input_chars: int = max_input_chars self.litellm_router_instance: Router = litellm_router_instance + self.encoder: LiteLLMRouterEncoder = LiteLLMRouterEncoder( + litellm_router_instance=litellm_router_instance, + model_name=embedding_model, + max_input_chars=max_input_chars, + ) def _load_semantic_routing_routes(self) -> list[Route]: from semantic_router.routers import SemanticRouter @@ -129,9 +154,6 @@ class AutoRouter(CustomLogger): from semantic_router.routers import SemanticRouter from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages - from litellm.router_strategy.auto_router.litellm_encoder import ( - LiteLLMRouterEncoder, - ) from litellm.types.router import PreRoutingHookResponse resolved_messages: Final = ( @@ -149,34 +171,47 @@ class AutoRouter(CustomLogger): ####################### routelayer = SemanticRouter( routes=self.loaded_routes, - encoder=LiteLLMRouterEncoder( - litellm_router_instance=self.litellm_router_instance, - model_name=self.embedding_model, - max_input_chars=self.max_input_chars, - ), + encoder=self.encoder, auto_sync=self.auto_sync_value, ) self.routelayer = routelayer message_content: Final = self._extract_text_from_messages(resolved_messages) - route_name: Final = self._matched_route_name(routelayer, message_content) + route_name: Final = await self._matched_route_name(routelayer, message_content, request_kwargs) return PreRoutingHookResponse( model=route_name or self.default_model, messages=messages, ) - def _matched_route_name(self, routelayer: "SemanticRouter", text: str) -> str | None: + async def _matched_route_name( + self, routelayer: "SemanticRouter", text: str, request_kwargs: Mapping[str, object] + ) -> str | None: """Name of the route `text` matches, or None when nothing matched or the match failed. - The route layer embeds `text` to compare it against the routes, and that embedding call can + `text` is embedded here rather than by `routelayer(text=...)` so the caller's metadata reaches + `aembedding()` and the embedding's spend lands on the key/team that sent the request; + SemanticRouter has no way to pass kwargs through to its encoder. That embedding call can fail (context limit, timeout, provider error). Choosing a model is a routing decision, so a failure here falls back to the default model rather than failing the user's request. """ from semantic_router.schema import RouteChoice try: - route_choice: Final = routelayer(text=text) + caller: Final = _CallerMetadata.model_validate(request_kwargs) + query_vector: Final = ( + await self.encoder.aencode_queries( + [text], + metadata=forwarded_internal_call_metadata(caller.metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN), + litellm_metadata=forwarded_internal_call_metadata( + caller.litellm_metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN + ), + proxy_server_request={"body": {"model": self.embedding_model, "input": [text]}}, + turn_off_message_logging=effective_turn_off_message_logging(request_kwargs), + **parent_session_kwargs(request_kwargs), + ) + )[0] + route_choice: Final = await routelayer.acall(vector=query_vector) except Exception as e: # noqa: BLE001 -- the embedding call behind the route layer can fail many ways (context limit, timeout, provider/network error); none of them may fail the request verbose_router_logger.warning( "AutoRouter: semantic routing failed (%s), falling back to default model %s", e, self.default_model diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_utils/auto_router_model_naming.py index a8aa543d735..2efbfb5782e 100644 --- a/litellm/router_utils/auto_router_model_naming.py +++ b/litellm/router_utils/auto_router_model_naming.py @@ -10,7 +10,7 @@ the router silently dropping the deployment at load time under ``ignore_invalid_deployments``. """ -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final, Literal, TypeAlias @@ -163,6 +163,38 @@ def strategy_router_dependencies( ) +def uses_heuristic_v2_classifier(complexity_router_config: object) -> bool: + """Whether this complexity config classifies with the bundled heuristic_v2 model.""" + return _mapping(complexity_router_config).get("classifier_type") == "heuristic_v2" + + +def is_heuristic_v2_router(litellm_params: Mapping[str, object]) -> bool: + """Whether this deployment is a complexity router that classifies with heuristic_v2.""" + return classify_strategy_router_model(str(litellm_params.get("model") or "")) == "complexity" and ( + uses_heuristic_v2_classifier(litellm_params.get("complexity_router_config")) + ) + + +def count_heuristic_v2_routers(deployments: Iterable[Mapping[str, object]]) -> int: + """How many of ``deployments`` (router model_list entries or config.yaml rows) are heuristic_v2 routers.""" + return sum(1 for deployment in deployments if is_heuristic_v2_router(_mapping(deployment.get("litellm_params")))) + + +def heuristic_v2_limit_violation(*, held: int, limit: int | None) -> str | None: + """Why holding ``held`` heuristic_v2 routers exceeds ``limit``, or None when it fits. + + ``limit`` None means unlimited. The message is shared by every enforcement point (config + load, model writes, router registration) and stays SDK-neutral: it names the cap and what + the caller can change; the proxy appends how its license lifts the cap. + """ + if limit is None or held <= limit: + return None + return ( + f"At most {limit} auto-router(s) with classifier_type 'heuristic_v2' can be registered but this would make " + f"{held}. Use classifier_type 'heuristic' for this router or remove an existing heuristic_v2 router." + ) + + def validate_complexity_router_config_write(complexity_router_config: Mapping[str, object] | None) -> str | None: """Reject a complexity config the router would refuse to build a deployment from. diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 39d3e25aacb..6f3ea8eb78a 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -82,6 +82,7 @@ class DeploymentAffinityCheck(CustomLogger): """ CACHE_KEY_PREFIX = "deployment_affinity:v1" + USER_ID_AFFINITY_PREFIX: Final = "user_id:" def __init__( self, @@ -253,15 +254,6 @@ class DeploymentAffinityCheck(CustomLogger): hashed_user_key: Final = cls._hash_user_key(user_key) if user_key is not None else "unscoped" return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}" - @staticmethod - def _get_user_key_from_metadata_dict(metadata: dict) -> str | None: - # NOTE: affinity is keyed on the *API key hash* provided by the proxy (not the - # OpenAI `user` parameter, which is an end-user identifier). - user_key: Final = metadata.get("user_api_key_hash") - if user_key is None: - return None - return str(user_key) - @staticmethod def _get_session_id_from_metadata_dict(metadata: dict) -> str | None: session_id: Final = metadata.get("session_id") @@ -285,22 +277,30 @@ class DeploymentAffinityCheck(CustomLogger): return metadata_dicts @staticmethod - def _get_user_key_from_request_kwargs(request_kwargs: dict) -> str | None: + def _first_metadata_value(metadata_dicts: Sequence[dict], key: str) -> str | None: + value: Final = next((metadata[key] for metadata in metadata_dicts if metadata.get(key) is not None), None) + return None if value is None else str(value) + + @classmethod + def _get_user_key_from_request_kwargs(cls, request_kwargs: dict) -> str | None: """ Extract a stable affinity key from request kwargs. - Source (proxy): `metadata.user_api_key_hash` + Source (proxy): `metadata.user_api_key_hash` for virtual-key callers. JWT-authenticated + callers carry no key hash, so their `metadata.user_api_key_user_id` stands in for it, + namespaced under `USER_ID_AFFINITY_PREFIX` so a user id can never alias a key hash. Note: the OpenAI `user` parameter is an end-user identifier and is intentionally not used for deployment affinity. """ - # Check metadata dicts (Proxy usage) - for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs): - user_key = DeploymentAffinityCheck._get_user_key_from_metadata_dict(metadata=metadata) - if user_key is not None: - return user_key - - return None + metadata_dicts: Final = cls._iter_metadata_dicts(request_kwargs) + user_api_key_hash: Final = cls._first_metadata_value(metadata_dicts, "user_api_key_hash") + if user_api_key_hash is not None: + return user_api_key_hash + user_id: Final = cls._first_metadata_value(metadata_dicts, "user_api_key_user_id") + if user_id is None: + return None + return f"{cls.USER_ID_AFFINITY_PREFIX}{user_id}" @staticmethod def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None: @@ -533,9 +533,9 @@ class DeploymentAffinityCheck(CustomLogger): return typed_healthy_deployments verbose_router_logger.debug( - "DeploymentAffinityCheck: api-key affinity hit -> deployment=%s user_key=%s", + "DeploymentAffinityCheck: caller affinity hit -> deployment=%s user_key=%s", model_id, - self._shorten_for_logs(user_key), + self._shorten_for_logs(self._hash_user_key(user_key)), ) return [deployment] @@ -626,7 +626,7 @@ class DeploymentAffinityCheck(CustomLogger): deployment_model_name, model_id, self.ttl_seconds, - self._shorten_for_logs(user_key), + self._shorten_for_logs(self._hash_user_key(user_key)), ) else: verbose_router_logger.debug( diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index 9a714e1724e..5de58a20242 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final from pydantic import BaseModel, Field @@ -29,6 +30,13 @@ def is_interception_internal_key( return any(key.startswith(prefix) for prefix in prefixes) +CONVERTED_STREAM_KEYS: Final = frozenset(f"{prefix}_converted_stream" for prefix in INTERCEPTION_INTERNAL_PREFIXES) + + +def converted_stream_requested(params: Mapping[str, object]) -> bool: + return any(bool(params.get(key)) for key in CONVERTED_STREAM_KEYS) + + class AgenticLoopSafetyError(ValueError): """ Raised when an agentic-loop safety rail refuses a rerun. diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 32d88da0085..b33ff954c35 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -66,6 +66,7 @@ from pydantic import ( ConfigDict, Discriminator, Field, + NonNegativeInt, PrivateAttr, SerializerFunctionWrapHandler, field_serializer, @@ -1321,6 +1322,18 @@ class ResponseAPIUsage(BaseLiteLLMOpenAIResponseObject): model_config = {"extra": "allow"} +class WebSearchToolUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + num_requests: NonNegativeInt + + +class ResponsesToolUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + web_search: WebSearchToolUsage | None = None + + ResponsesAPIStatus = Literal["completed", "failed", "in_progress", "cancelled", "queued", "incomplete"] """ The status of the response generation. diff --git a/litellm/types/proxy/model_listing.py b/litellm/types/proxy/model_listing.py index b59c0f2cf19..24cfa85eee4 100644 --- a/litellm/types/proxy/model_listing.py +++ b/litellm/types/proxy/model_listing.py @@ -11,8 +11,8 @@ class ModelInfoMetadata(TypedDict): class ModelInfoResponse(TypedDict): """OpenAI-compatible model object. `mode`, `max_input_tokens`, and - `max_output_tokens` are attached when the cost map knows them; `metadata` - is present only when the endpoint is called with include_metadata=true. + `max_output_tokens` are attached when the cost map or deployment config + knows them; `metadata` is present only with include_metadata=true. """ id: str diff --git a/litellm/types/router.py b/litellm/types/router.py index 4f4df1a8d2e..7ebd50f1328 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -885,6 +885,18 @@ class FallbackAccessCheck(Protocol): async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ... +class HeuristicV2RouterLimit(Protocol): + """ + Resolves how many heuristic_v2 complexity routers the Router may hold right now; None means unlimited. + + The Router calls it on every registration and limit query instead of caching the answer, so the + proxy can keep the limit on its license object (re-verified on config load) rather than hand + over a snapshot. + """ + + def __call__(self) -> int | None: ... + + class LiteLLM_RouterFileObject(TypedDict, total=False): """ Tracking the litellm params hash, used for mapping the file id to the right model diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3b6e7912253..1c51385318b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29354,6 +29354,75 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, + "gpt-6-astra": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, + "cache_creation_input_token_cost_flex": 6.25e-06, + "cache_creation_input_token_cost_priority": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 4e-06, + "cache_read_input_token_cost_flex": 5e-07, + "cache_read_input_token_cost_priority": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "input_cost_per_token_above_272k_tokens_flex": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 4e-05, + "input_cost_per_token_batches": 5e-06, + "input_cost_per_token_flex": 5e-06, + "input_cost_per_token_priority": 2e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, + "output_cost_per_token_above_272k_tokens_priority": 0.00015, + "output_cost_per_token_batches": 2.5e-05, + "output_cost_per_token_flex": 2.5e-05, + "output_cost_per_token_priority": 0.0001, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "gpt-5.6": { "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, @@ -53485,6 +53554,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53519,6 +53593,11 @@ "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, "output_cost_per_token": 1.32e-05, "output_cost_per_token_above_272k_tokens": 1.98e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53581,6 +53660,11 @@ "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, "output_cost_per_token": 1.32e-06, "output_cost_per_token_above_272k_tokens": 1.98e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53769,6 +53853,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53800,6 +53889,11 @@ "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "output_cost_per_token": 1.65e-05, "output_cost_per_token_above_272k_tokens": 2.475e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index ef8d6c55148..f8f23ea015a 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -2879,7 +2879,6 @@ def response_format_tests(response: litellm.ModelResponse): "model", [ "bedrock/mistral.mistral-large-2407-v1:0", - "bedrock/cohere.command-r-plus-v1:0", "us.anthropic.claude-sonnet-4-5-20250929-v1:0", "mistral.mistral-7b-instruct-v0:2", "meta.llama3-8b-instruct-v1:0", diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index 07d693af447..bf39d3155b7 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -1168,7 +1168,6 @@ async def test_completion_replicate_llama3_streaming(sync_mode): "model, region", [ # ["bedrock/ai21.jamba-instruct-v1:0", "us-east-1"], - # ["bedrock/cohere.command-r-plus-v1:0", None], ["us.anthropic.claude-sonnet-4-5-20250929-v1:0", None], # ["mistral.mistral-7b-instruct-v0:2", None], # ["meta.llama3-8b-instruct-v1:0", None], @@ -1271,7 +1270,7 @@ def test_bedrock_claude_3_streaming(): "model", [ "claude-haiku-4-5-20251001", - "cohere.command-r-plus-v1:0", # bedrock + "bedrock/mistral.mistral-7b-instruct-v0:2", "gpt-3.5-turbo", ], ) diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 21b60d7a216..6590718878d 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -1295,6 +1295,19 @@ def test_text_plus_tool_calls_sequence(): # ============================================================================= +def test_developer_message_content_uses_input_text(): + handler = LiteLLMResponsesTransformationHandler() + + input_items, instructions = handler.convert_chat_completion_messages_to_responses_api( + [{"role": "developer", "content": "Always answer in French."}] + ) + + assert instructions is None + assert input_items == [ + {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "Always answer in French."}]} + ] + + def test_tool_message_output_uses_input_text_not_output_text(): """ Test that tool message content uses input_text type, not output_text. diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index b7f0ca1efe1..0b6832d4bef 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1662,6 +1662,54 @@ def test_generic_cost_per_token_gpt56_cyber( assert completion_cost == pytest.approx(completion_tokens * output_rate) +@pytest.mark.parametrize( + "service_tier,tier_multiplier", + [(None, 1.0), ("flex", 0.5), ("priority", 2.0), ("fast", 2.0)], +) +@pytest.mark.parametrize( + "prompt_tokens,input_side_multiplier,output_multiplier", + [(100000, 1.0, 1.0), (300000, 2.0, 1.5)], +) +def test_generic_cost_per_token_gpt_6_astra_price_sheet( + _local_model_cost_map, + service_tier, + tier_multiplier, + prompt_tokens, + input_side_multiplier, + output_multiplier, +): + """gpt-6-astra launch price sheet: $10 input, $1 cache read, $12.50 cache write, $50 output per 1M tokens. + + Above 272K prompt tokens the input-side rates double and the output rate is 1.5x on the whole + request. Flex is half the applicable rate and fast mode, billed as priority, is double it. + """ + cached_tokens = 50000 + cache_write_tokens = 40000 + text_tokens = prompt_tokens - cached_tokens - cache_write_tokens + completion_tokens = 1000 + usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=cached_tokens, cache_write_tokens=cache_write_tokens + ), + ) + + prompt_cost, completion_cost = generic_cost_per_token( + model="gpt-6-astra", + usage=usage, + custom_llm_provider="openai", + service_tier=service_tier, + ) + + input_side = tier_multiplier * input_side_multiplier + assert prompt_cost == pytest.approx( + input_side * (text_tokens * 1e-5 + cached_tokens * 1e-6 + cache_write_tokens * 1.25e-5) + ) + assert completion_cost == pytest.approx(tier_multiplier * output_multiplier * completion_tokens * 5e-5) + + @pytest.mark.parametrize( "model,input_cost,output_cost,cache_read_cost", [ diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index fd795ffcc96..6cc3dcceebc 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,4 +1,5 @@ import os +from collections.abc import Mapping, Sequence import pytest @@ -6,7 +7,7 @@ import litellm from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) -from litellm.types.llms.openai import FileSearchTool, WebSearchOptions +from litellm.types.llms.openai import FileSearchTool, ResponsesAPIResponse, WebSearchOptions from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams @@ -928,3 +929,125 @@ def test_web_search_gate_reads_server_side_tool_usage_details_without_citations( standard_built_in_tools_params=None, ) assert cost == 3 * _DEFAULT_WEB_SEARCH_COST_PER_CALL + + +_BEDROCK_MANTLE_WEB_SEARCH_MODELS = ( + "bedrock_mantle/openai.gpt-5.6-sol", + "bedrock_mantle/openai.gpt-5.6-terra", + "bedrock_mantle/openai.gpt-5.6-luna", + "bedrock_mantle/openai.gpt-5.5", + "bedrock_mantle/openai.gpt-5.4", +) + +_BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012 + + +def _responses_with_web_search( + model: str, actions: Sequence[Mapping[str, str]], tool_usage: Mapping[str, object] | None = None +) -> ResponsesAPIResponse: + payload = { + "id": "resp_1", + "created_at": 1756900000, + "model": model.split("/", 1)[-1], + "object": "response", + "status": "completed", + "output": [ + {"type": "web_search_call", "id": f"ws_{i}", "status": "completed", "action": action} + for i, action in enumerate(actions) + ], + } + return ResponsesAPIResponse.model_validate( + payload if tool_usage is None else {**payload, "tool_usage": tool_usage} + ) + + +def _web_search_cost(model: str, response: ResponsesAPIResponse, custom_llm_provider: str) -> float: + from litellm.types.utils import Usage + + return StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=response, + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + custom_llm_provider=custom_llm_provider, + standard_built_in_tools_params=None, + ) + + +@pytest.mark.parametrize("model", _BEDROCK_MANTLE_WEB_SEARCH_MODELS) +def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model): + """Two Bedrock-reported web searches bill 2 x $0.012 under the prefixed and the bare model id alike.""" + pricing = litellm.get_model_info(model)["search_context_cost_per_query"] + assert pricing == { + "search_context_size_low": _BEDROCK_MANTLE_WEB_SEARCH_RATE, + "search_context_size_medium": _BEDROCK_MANTLE_WEB_SEARCH_RATE, + "search_context_size_high": _BEDROCK_MANTLE_WEB_SEARCH_RATE, + } + + response = _responses_with_web_search( + model, + actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}], + tool_usage={"web_search": {"num_requests": 2}}, + ) + for cost_model in (model, model.split("/", 1)[1]): + cost = _web_search_cost(cost_model, response, "bedrock_mantle") + assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( + f"{cost_model} must bill 2 x ${_BEDROCK_MANTLE_WEB_SEARCH_RATE} for 2 web searches, got ${cost}" + ) + + +@pytest.mark.parametrize("num_requests", [1, 0]) +def test_web_search_call_count_prefers_provider_reported_num_requests(local_model_cost_map, num_requests): + """A search plus an open_page fetch bills tool_usage.web_search.num_requests, never the two items.""" + model = "bedrock_mantle/openai.gpt-5.6-sol" + response = _responses_with_web_search( + model, + actions=[ + {"type": "search", "query": "litellm"}, + {"type": "open_page", "url": "https://docs.litellm.ai/"}, + ], + tool_usage={"web_search": {"num_requests": num_requests}}, + ) + + cost = _web_search_cost(model, response, "bedrock_mantle") + + assert cost == pytest.approx(num_requests * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( + f"{num_requests} reported web search requests must bill {num_requests} x " + f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}" + ) + + +@pytest.mark.parametrize( + "tool_usage", + [None, {}, {"web_search": None}, {"web_search": {"num_requests": "many"}}, {"web_search": {"num_requests": -1}}], +) +def test_web_search_call_count_falls_back_to_items_without_reported_count(local_model_cost_map, tool_usage): + """Without a usable reported count the per-call path keeps counting web_search_call items.""" + model = "bedrock_mantle/openai.gpt-5.6-sol" + response = _responses_with_web_search( + model, + actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}], + tool_usage=tool_usage, + ) + + cost = _web_search_cost(model, response, "bedrock_mantle") + + assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( + f"2 web_search_call items with tool_usage={tool_usage!r} must bill 2 x " + f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}" + ) + + +def test_web_search_call_count_reads_reported_count_beside_other_tool_usage_entries(local_model_cost_map): + """OpenAI reports web_search.num_requests next to other tool entries, which must not disable the reported count.""" + response = _responses_with_web_search( + "gpt-5.6", + actions=[{"type": "search", "query": "S&P 500 close"}, {"type": "open_page", "url": "https://example.com/"}], + tool_usage={ + "image_gen": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}, + "web_search": {"num_requests": 1}, + }, + ) + + cost = _web_search_cost("gpt-5.6", response, "openai") + + assert cost == pytest.approx(0.01), f"1 reported OpenAI web search must bill 1 x $0.01, not the 2 items, got ${cost}" diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 4b2a3f9247c..0aa73833677 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4761,6 +4761,43 @@ async def test_async_stream_assembled_response_keeps_vertex_traffic_type(logging assert assembled._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND_FLEX" +@pytest.mark.asyncio +async def test_async_fake_stream_final_chunk_carries_hidden_usage(logging_obj: Logging): + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator + from litellm.types.utils import ModelResponse + + model_response = ModelResponse( + id="chatcmpl-fake-stream", + model="my-random-model", + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "hello world"}, + "finish_reason": "stop", + } + ], + ) + model_response.usage = Usage(prompt_tokens=1234, completion_tokens=7, total_tokens=1241) + + wrapper = CustomStreamWrapper( + completion_stream=MockResponseIterator(model_response=model_response), + model="my-random-model", + custom_llm_provider="anthropic", + logging_obj=logging_obj, + ) + + final_chunk = None + async for chunk in wrapper: + final_chunk = chunk + + assert final_chunk is not None + hidden_usage = final_chunk._hidden_params.get("usage") + assert hidden_usage is not None + assert hidden_usage.prompt_tokens == 1234 + assert hidden_usage.completion_tokens == 7 + assert hidden_usage.total_tokens == 1241 + + class TestStableStreamingResponseId: """ All chunks of one streamed response must share the same top-level id diff --git a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py index e06cae97283..bd0f16a695b 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py @@ -336,3 +336,15 @@ class TestAzureResolvesTheDeclaredDefaultEffort: drop_params=True, ) assert ("temperature" in mapped) is temperature_survives + + +def test_azure_gpt_6_astra_takes_the_reasoning_series_request_shape(): + params = litellm.get_optional_params( + model="gpt-6-astra", + custom_llm_provider="azure", + max_tokens=100, + reasoning_effort="max", + ) + assert params["max_completion_tokens"] == 100 + assert "max_tokens" not in params + assert params["reasoning_effort"] == "max" diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 11a727c9635..33fbb4e8fc7 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -31,6 +31,46 @@ async def test_get_openai_compatible_provider_info(): assert custom_llm_provider == "azure" +@pytest.mark.parametrize( + "model, api_base, expected_provider", + [ + ("azure_ai/gpt-4o", "https://my-resource.services.ai.azure.com", "azure_ai"), + ("azure_ai/gpt-4o", "https://my-resource.services.ai.azure.com/models", "azure_ai"), + ("azure_ai/gpt-5.4-nano", "https://my-resource.services.ai.azure.com", "azure_ai"), + ("azure_ai/gpt-4o", "https://my-resource.openai.azure.com", "azure"), + ( + "azure_ai/gpt-4o", + "https://my-resource.services.ai.azure.com/openai/deployments/gpt-4o/chat/completions" + "?api-version=2024-08-01-preview", + "azure", + ), + ("azure_ai/mistral-large-latest", "https://my-resource.services.ai.azure.com", "azure_ai"), + ("azure_ai/mistral-large-latest", "https://my-resource.openai.azure.com", "azure_ai"), + ], +) +def test_foundry_base_keeps_azure_ai_provider(model: str, api_base: str, expected_provider: str): + """Regression for #38276: a Foundry .services.ai.azure.com base must not be reclassified as azure.""" + config = AzureAIStudioConfig() + ( + _, + _, + custom_llm_provider, + ) = config._get_openai_compatible_provider_info( + model=model, + api_base=api_base, + api_key="my-key", + custom_llm_provider="azure_ai", + ) + assert custom_llm_provider == expected_provider + + +def test_is_azure_openai_model_without_api_base_keeps_azure_ai(): + """Metadata lookups (get_model_info, supports_* checks) carry no api_base and must not flip the provider.""" + config = AzureAIStudioConfig() + assert config._is_azure_openai_model(model="azure_ai/gpt-4o", api_base=None) is False + assert config._is_azure_openai_model(model="azure_ai/gpt-4o", api_base="https://my-res.openai.azure.com") is True + + def test_azure_ai_validate_environment(): config = AzureAIStudioConfig() headers = config.validate_environment( diff --git a/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py b/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py new file mode 100644 index 00000000000..0401629171d --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py @@ -0,0 +1,69 @@ +import httpx +import pytest +import respx + +from litellm import embedding +from litellm.llms.azure_ai.embed.handler import _foundry_models_route_base + +EMBEDDING_PAYLOAD = { + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, +} + + +@pytest.mark.parametrize( + ("api_base", "expected"), + [ + ( + "https://my-foundry.services.ai.azure.com", + "https://my-foundry.services.ai.azure.com/models", + ), + ( + "https://my-foundry.services.ai.azure.com/", + "https://my-foundry.services.ai.azure.com/models", + ), + ( + "https://my-foundry.services.ai.azure.com?api-version=2024-05-01-preview", + "https://my-foundry.services.ai.azure.com/models?api-version=2024-05-01-preview", + ), + ( + "https://my-foundry.services.ai.azure.com/models", + "https://my-foundry.services.ai.azure.com/models", + ), + ( + "https://my-foundry.services.ai.azure.com/openai/deployments/text-embedding-3-small", + "https://my-foundry.services.ai.azure.com/openai/deployments/text-embedding-3-small", + ), + ( + "https://my-resource.openai.azure.com", + "https://my-resource.openai.azure.com", + ), + ( + "https://Mistral-serverless.eastus2.models.ai.azure.com", + "https://Mistral-serverless.eastus2.models.ai.azure.com", + ), + (None, None), + ], +) +def test_foundry_models_route_base(api_base, expected): + assert _foundry_models_route_base(api_base) == expected + + +@respx.mock +def test_azure_ai_embedding_calls_foundry_models_route(): + route = respx.post("https://my-foundry.services.ai.azure.com/models/embeddings").mock( + return_value=httpx.Response(200, json=EMBEDDING_PAYLOAD) + ) + + response = embedding( + model="azure_ai/text-embedding-3-small", + input=["hello world"], + api_base="https://my-foundry.services.ai.azure.com", + api_key="fake-key", + ) + + assert route.called + assert response.data is not None + assert len(response.data) == 1 diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py index 8dc4620dd1b..b6281834f24 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py @@ -10,11 +10,11 @@ Tests the cost calculation for Dashscope models including: import math import os +from datetime import datetime, timezone import pytest # Add the project root to Python path - import litellm from litellm.llms.dashscope.cost_calculator import ( cost_per_token as dashscope_cost_per_token, @@ -526,3 +526,139 @@ class TestDashscopeCostCalculator: assert prompt_cost == 0.0 assert math.isclose(completion_cost, 500 * 1.6e-06, rel_tol=1e-10) + + OFF_PEAK_WINDOW = "14:00-00:00" + INSIDE_WINDOW = datetime(2026, 9, 3, 17, 25, tzinfo=timezone.utc) + OUTSIDE_WINDOW = datetime(2026, 9, 3, 9, 0, tzinfo=timezone.utc) + + def _register_off_peak_flat_model(self, model_key: str, off_peak_pricing: dict) -> None: + litellm.model_cost[model_key] = { + "litellm_provider": "dashscope", + "mode": "chat", + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 4.8e-06, + "cache_read_input_token_cost": 2e-07, + "cache_creation_input_token_cost": 3e-06, + "off_peak_pricing": off_peak_pricing, + } + + def test_dashscope_off_peak_window_swaps_in_the_off_peak_rates(self): + """ + Regression (LIT-6782): a deployment configured with off_peak_pricing kept billing the + standard dashscope rates inside its window, while the same block on a deepseek + deployment billed the off-peak rates. + """ + self._register_off_peak_flat_model( + "dashscope/deepseek-off-peak-test", + { + "hours_utc": self.OFF_PEAK_WINDOW, + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 2.4e-06, + "cache_read_input_token_cost": 1e-07, + }, + ) + usage = Usage( + prompt_tokens=1000, + completion_tokens=200, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300, cache_creation_tokens=100), + ) + + prompt_cost, completion_cost = dashscope_cost_per_token( + model="deepseek-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + + assert math.isclose(prompt_cost, (600 * 1.2e-06) + (300 * 1e-07) + (100 * 3e-06), rel_tol=1e-10) + assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10) + + peak_prompt_cost, peak_completion_cost = dashscope_cost_per_token( + model="deepseek-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW + ) + + assert math.isclose(peak_prompt_cost, (600 * 2.4e-06) + (300 * 2e-07) + (100 * 3e-06), rel_tol=1e-10) + assert math.isclose(peak_completion_cost, 200 * 4.8e-06, rel_tol=1e-10) + + def test_dashscope_off_peak_window_overrides_the_selected_tier(self): + """An open off-peak window bills the whole request at the flat off-peak rates, whichever tier + the input volume selected.""" + self._register_tiered_model( + "dashscope/qwen-tiered-off-peak-test", + [ + {"range": [0, 1000], "input_cost_per_token": 4e-07, "output_cost_per_token": 1.6e-06}, + {"range": [1000, 2000], "input_cost_per_token": 8e-07, "output_cost_per_token": 3.2e-06}, + ], + ) + litellm.model_cost["dashscope/qwen-tiered-off-peak-test"]["off_peak_pricing"] = { + "hours_utc": self.OFF_PEAK_WINDOW, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 4e-07, + } + usage = Usage(prompt_tokens=1500, completion_tokens=300) + + prompt_cost, completion_cost = dashscope_cost_per_token( + model="qwen-tiered-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + + assert math.isclose(prompt_cost, 1500 * 1e-07, rel_tol=1e-10) + assert math.isclose(completion_cost, 300 * 4e-07, rel_tol=1e-10) + + peak_prompt_cost, peak_completion_cost = dashscope_cost_per_token( + model="qwen-tiered-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW + ) + + assert math.isclose(peak_prompt_cost, 1500 * 8e-07, rel_tol=1e-10) + assert math.isclose(peak_completion_cost, 300 * 3.2e-06, rel_tol=1e-10) + + def test_dashscope_off_peak_rates_left_unset_keep_the_standard_rates(self): + """A block that only overrides the input rate leaves output and cache reads on the standard + rates, and an explicit reasoning rate is never swapped out.""" + self._register_off_peak_flat_model( + "dashscope/qwen-partial-off-peak-test", + {"hours_utc": self.OFF_PEAK_WINDOW, "input_cost_per_token": 1.2e-06}, + ) + litellm.model_cost["dashscope/qwen-partial-off-peak-test"]["output_cost_per_reasoning_token"] = 9e-06 + usage = Usage( + prompt_tokens=1000, + completion_tokens=200, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=50), + ) + + prompt_cost, completion_cost = dashscope_cost_per_token( + model="qwen-partial-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + + assert math.isclose(prompt_cost, (700 * 1.2e-06) + (300 * 2e-07), rel_tol=1e-10) + assert math.isclose(completion_cost, (150 * 4.8e-06) + (50 * 9e-06), rel_tol=1e-10) + + def test_dashscope_off_peak_output_rate_covers_reasoning_without_a_dedicated_rate(self): + """Reasoning tokens on a model with no dedicated reasoning rate follow the off-peak output + rate, the same way they follow the standard output rate outside the window.""" + self._register_off_peak_flat_model( + "dashscope/qwen-reasoning-off-peak-test", + {"hours_utc": self.OFF_PEAK_WINDOW, "output_cost_per_token": 2.4e-06}, + ) + usage = Usage( + prompt_tokens=100, + completion_tokens=200, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=50), + ) + + _, completion_cost = dashscope_cost_per_token( + model="qwen-reasoning-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + + assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10) + + def test_dashscope_off_peak_defaults_to_the_current_time(self): + """The proxy's cost dispatch passes no clock, so an all-day window has to apply on the + default current time.""" + self._register_off_peak_flat_model( + "dashscope/qwen-all-day-off-peak-test", + {"hours_utc": "00:00-00:00", "input_cost_per_token": 1.2e-06, "output_cost_per_token": 2.4e-06}, + ) + usage = Usage(prompt_tokens=1000, completion_tokens=200) + + prompt_cost, completion_cost = dashscope_cost_per_token(model="qwen-all-day-off-peak-test", usage=usage) + + assert math.isclose(prompt_cost, 1000 * 1.2e-06, rel_tol=1e-10) + assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 295121167d6..a453708040e 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1329,6 +1329,523 @@ class TestOpenAIResponsesHandlerToolInjection: assert "injected_tool" in names +COMPRESSED_MARKER = "[compressed document; retrieve the full text with hash=b573993006976af767214fac]" + + +class StructuredRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites whole messages via structured_messages and leaves + texts untouched, the way message-compressing guardrails do.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first_user = next(i for i, m in enumerate(messages) if m.get("role") == "user") + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten} + + +class ToolOutputRewriteGuardrail(CustomGuardrail): + """Guardrail that compresses the first tool-result row, the way Headroom does.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first_tool = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "tool") + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if i == first_tool else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten} + + +class DroppingRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites the first user row and drops the last row, so the + rewrite can only land through the full-conversion fallback.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first_user = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "user") + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten[:-1]} + + +def _texts(item: dict) -> list[str]: + content = item.get("content") + if isinstance(content, str): + return [content] + return [part["text"] for part in content] + + +class TestStructuredMessagesWriteBack: + """A guardrail's structured_messages rewrite must land in the Responses request, + not only the per-text mapping the chat handler shares with it.""" + + @pytest.mark.asyncio + async def test_list_input_gets_rewritten_messages_and_keeps_instructions(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [ + {"role": "user", "content": "memo " * 400}, + {"role": "assistant", "content": "Understood."}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert result["instructions"] == "Answer from the memo only." + user_items = [item for item in result["input"] if item.get("role") == "user"] + assert [_texts(item) for item in user_items] == [[COMPRESSED_MARKER], ["What is the codename?"]] + assert not any(item.get("role") == "system" for item in result["input"]) + assert _texts(next(item for item in result["input"] if item.get("role") == "assistant")) == ["Understood."] + + @pytest.mark.asyncio + async def test_string_input_becomes_rewritten_message_list(self): + handler = OpenAIResponsesHandler() + data = {"model": "gpt-5.6", "input": "memo " * 400} + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert [_texts(item) for item in result["input"]] == [[COMPRESSED_MARKER]] + assert "instructions" not in result + + @pytest.mark.asyncio + async def test_developer_item_preserved_verbatim_by_row_patch(self): + handler = OpenAIResponsesHandler() + developer_item = {"role": "developer", "content": "Always answer in French."} + data = { + "model": "gpt-5.6", + "input": [ + developer_item, + {"role": "user", "content": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert result["input"][0] is developer_item + assert developer_item["content"] == "Always answer in French." + assert _texts(result["input"][1]) == [COMPRESSED_MARKER] + assert _texts(result["input"][2]) == ["What is the codename?"] + + @pytest.mark.asyncio + async def test_reasoning_and_function_call_items_survive_tool_output_compression(self): + handler = OpenAIResponsesHandler() + reasoning_item = { + "id": "rs_123", + "type": "reasoning", + "summary": [], + "encrypted_content": "gAAAAA-signed-reasoning", + } + function_call_item = { + "id": "fc_123", + "type": "function_call", + "call_id": "call_abc", + "name": "read_document", + "arguments": '{"path": "memo.txt"}', + "status": "completed", + } + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [ + reasoning_item, + function_call_item, + {"type": "function_call_output", "call_id": "call_abc", "output": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, ToolOutputRewriteGuardrail()) + + assert result["instructions"] == "Answer from the memo only." + assert result["input"][0] is reasoning_item + assert reasoning_item["encrypted_content"] == "gAAAAA-signed-reasoning" + assert result["input"][1] is function_call_item + assert function_call_item["id"] == "fc_123" + assert result["input"][2] == { + "type": "function_call_output", + "call_id": "call_abc", + "output": COMPRESSED_MARKER, + } + assert result["input"][3] == {"role": "user", "content": "What is the codename?"} + + @pytest.mark.asyncio + async def test_web_search_call_item_preserved_verbatim(self): + handler = OpenAIResponsesHandler() + web_search_item = { + "id": "ws_123", + "type": "web_search_call", + "status": "completed", + "action": {"type": "search", "query": "codename memo"}, + } + data = { + "model": "gpt-5.6", + "input": [ + web_search_item, + {"role": "user", "content": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert result["input"][0] is web_search_item + assert _texts(result["input"][1]) == [COMPRESSED_MARKER] + assert _texts(result["input"][2]) == ["What is the codename?"] + + @pytest.mark.asyncio + async def test_row_count_change_falls_back_to_full_conversion(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "input": [ + {"role": "developer", "content": "Always answer in French."}, + {"role": "user", "content": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, DroppingRewriteGuardrail()) + + assert len(result["input"]) == 2 + developer = next(item for item in result["input"] if item.get("role") == "developer") + assert developer["content"] == [{"type": "input_text", "text": "Always answer in French."}] + assert _texts(next(item for item in result["input"] if item.get("role") == "user")) == [COMPRESSED_MARKER] + + @pytest.mark.asyncio + async def test_same_inputs_object_back_keeps_the_text_mapping(self): + handler = OpenAIResponsesHandler() + original_input = [ + {"role": "user", "content": "Hello"}, + {"role": "user", "content": [{"type": "input_text", "text": "Again"}]}, + ] + data = {"model": "gpt-5.6", "input": original_input} + + result = await handler.process_input_messages(data, MockGuardrail()) + + assert result["input"] is original_input + assert [_texts(item) for item in result["input"]] == [["Hello [GUARDRAILED]"], ["Again [GUARDRAILED]"]] + + +class AllToolOutputsRewriteGuardrail(CustomGuardrail): + """Guardrail that compresses every tool-result row.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if isinstance(m, dict) and m.get("role") == "tool" else m + for m in messages + ] + return {**inputs, "structured_messages": rewritten} + + +class AssistantRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites the first assistant row's content.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "assistant") + rewritten = [{**m, "content": COMPRESSED_MARKER} if i == first else m for i, m in enumerate(messages)] + return {**inputs, "structured_messages": rewritten} + + +class DictStructuredMessagesGuardrail(CustomGuardrail): + """Guardrail that hands back a raw evaluation dict instead of a message list, + the way HiddenLayer v2 does.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "structured_messages": {"evaluation": "allowed", "messages": []}} + + +def _parallel_tool_call_input() -> list: + return [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"id": "fc_2", "type": "function_call", "call_id": "call_2", "name": "read_b", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": "memo " * 400}, + {"type": "function_call_output", "call_id": "call_2", "output": "note " * 400}, + {"role": "user", "content": "What is the codename?"}, + ] + + +class TestProvenancePatching: + """The O(n) provenance pass must keep patching rewritten rows in place for the + shapes real agent loops produce, and fall back safely everywhere else.""" + + @pytest.mark.asyncio + async def test_parallel_tool_call_outputs_both_patched(self): + handler = OpenAIResponsesHandler() + raw_input = _parallel_tool_call_input() + fc_1, fc_2 = raw_input[0], raw_input[1] + data = {"model": "gpt-5.6", "input": raw_input} + + result = await handler.process_input_messages(data, AllToolOutputsRewriteGuardrail()) + + assert result["input"][0] is fc_1 + assert result["input"][1] is fc_2 + assert result["input"][2] == {"type": "function_call_output", "call_id": "call_1", "output": COMPRESSED_MARKER} + assert result["input"][3] == {"type": "function_call_output", "call_id": "call_2", "output": COMPRESSED_MARKER} + assert result["input"][4] == {"role": "user", "content": "What is the codename?"} + + @pytest.mark.asyncio + async def test_assistant_turn_with_tool_call_keeps_items_verbatim(self): + handler = OpenAIResponsesHandler() + assistant_item = {"role": "assistant", "content": "Let me read the memo."} + function_call_item = { + "id": "fc_9", + "type": "function_call", + "call_id": "call_9", + "name": "read_document", + "arguments": '{"path": "memo.txt"}', + } + data = { + "model": "gpt-5.6", + "input": [ + assistant_item, + function_call_item, + {"type": "function_call_output", "call_id": "call_9", "output": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, ToolOutputRewriteGuardrail()) + + assert result["input"][0] is assistant_item + assert result["input"][1] is function_call_item + assert result["input"][2] == {"type": "function_call_output", "call_id": "call_9", "output": COMPRESSED_MARKER} + + @pytest.mark.asyncio + async def test_rewrite_of_merged_tool_call_message_falls_back(self): + handler = OpenAIResponsesHandler() + raw_input = _parallel_tool_call_input() + data = {"model": "gpt-5.6", "input": raw_input} + + result = await handler.process_input_messages(data, AssistantRewriteGuardrail()) + + assert not any(item is original for item in result["input"] for original in raw_input) + assistant_items = [item for item in result["input"] if item.get("role") == "assistant"] + assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]] + + @pytest.mark.asyncio + async def test_rewrite_of_lone_function_call_message_falls_back(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "input": [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": "memo memo"}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + raw_input = data["input"] + result = await handler.process_input_messages(data, AssistantRewriteGuardrail()) + + assert not any(item is original for item in result["input"] for original in raw_input) + assistant_items = [item for item in result["input"] if item.get("role") == "assistant"] + assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]] + + def test_provenance_bails_on_non_mapping_item(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + + assert _input_item_provenance(["not a mapping"], []) is None + + def test_provenance_bails_when_expected_messages_disagree(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + + assert _input_item_provenance([{"role": "user", "content": "hi"}], [{"role": "user", "content": "bye"}]) is None + + def test_provenance_bails_on_unpredicted_merge(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + raw_input = [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"role": "assistant", "content": "Reading the memo now."}, + ] + expected = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=raw_input, responses_api_request={} + ) + assert len(expected) == 1 + assert _input_item_provenance(raw_input, expected) is None + + def test_provenance_maps_and_taints_parallel_tool_calls(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + raw_input = _parallel_tool_call_input() + expected = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=raw_input, responses_api_request={} + ) + provenance = _input_item_provenance(raw_input, expected) + assert provenance is not None + item_for_message, tainted = provenance + assert tainted == {0} + assert dict(item_for_message) == {1: 2, 2: 3, 3: 4} + + +class TestDictStructuredMessagesGuard: + """A guardrail handing back a non-list structured_messages payload must not + blow up the request; the write-back is skipped instead.""" + + @pytest.mark.asyncio + async def test_list_input_survives_dict_structured_messages(self): + handler = OpenAIResponsesHandler() + original_input = [{"role": "user", "content": "Hello"}] + data = {"model": "gpt-5.6", "input": original_input} + + result = await handler.process_input_messages(data, DictStructuredMessagesGuardrail()) + + assert result["input"] is original_input + assert result["input"] == [{"role": "user", "content": "Hello"}] + + @pytest.mark.asyncio + async def test_string_input_survives_dict_structured_messages(self): + handler = OpenAIResponsesHandler() + data = {"model": "gpt-5.6", "input": "Hello there"} + + result = await handler.process_input_messages(data, DictStructuredMessagesGuardrail()) + + assert result["input"] == "Hello there" + + +class SystemRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites the system row, the way prompt-hardening guardrails do.""" + + def __init__(self, rewritten_content: Any = COMPRESSED_MARKER): + super().__init__() + self.rewritten_content = rewritten_content + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "system") + rewritten = [ + {**m, "content": self.rewritten_content} if i == first else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten} + + +class TestPatchEdgeBranches: + @pytest.mark.asyncio + async def test_multimodal_user_item_rewritten_through_conversion(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "input": [ + {"role": "user", "content": [{"type": "input_text", "text": "memo " * 400}]}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert _texts(result["input"][0]) == [COMPRESSED_MARKER] + assert result["input"][1] == {"role": "user", "content": "What is the codename?"} + + @pytest.mark.asyncio + async def test_instructions_rewrite_lands_in_instructions_field(self): + handler = OpenAIResponsesHandler() + user_item = {"role": "user", "content": "What is the codename?"} + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [user_item], + } + + result = await handler.process_input_messages(data, SystemRewriteGuardrail()) + + assert result["instructions"] == COMPRESSED_MARKER + assert result["input"][0] is user_item + + @pytest.mark.asyncio + async def test_non_string_instructions_rewrite_falls_back(self): + handler = OpenAIResponsesHandler() + user_item = {"role": "user", "content": "What is the codename?"} + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [user_item], + } + + result = await handler.process_input_messages( + data, SystemRewriteGuardrail(rewritten_content=[{"type": "text", "text": COMPRESSED_MARKER}]) + ) + + assert result["input"][0] is not user_item + + @pytest.mark.asyncio + async def test_unpredicted_merge_falls_back_through_patch(self): + handler = OpenAIResponsesHandler() + raw_input = [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"role": "assistant", "content": "Reading the memo now."}, + {"type": "function_call_output", "call_id": "call_1", "output": "memo memo"}, + {"role": "user", "content": "memo " * 400}, + ] + data = {"model": "gpt-5.6", "input": raw_input} + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert not any(item is original for item in result["input"] for original in raw_input) + user_items = [item for item in result["input"] if item.get("role") == "user"] + assert _texts(user_items[0]) == [COMPRESSED_MARKER] + + def test_item_rewrite_field_ignores_non_string_type(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _item_rewrite_field + + assert _item_rewrite_field({"type": 123, "content": "hello"}) is None + + class ToolEditingGuardrail(CustomGuardrail): """Guardrail that rewrites the flattened chat tools it was handed through ``edit``""" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index b0ffd1845fe..4ac072d0ca6 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -1718,6 +1718,8 @@ class TestResponsesSurfaceSharesTheEffortRule: ("gpt-5.6-sol", None, False), ("gpt-5.6-terra", "none", True), ("gpt-5.6-terra", "medium", False), + ("gpt-6-astra", None, False), + ("gpt-6-astra", "low", False), ], ) def test_temperature_follows_the_resolved_effort( diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index 9c5bd34d59a..c86ce4df2ac 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -1505,3 +1505,17 @@ class TestACatalogueOlderThanTheCodeDoesNotStripTemperature: drop_params=True, ) assert "temperature" not in mapped + + +def test_gpt_6_astra_takes_the_reasoning_series_request_shape(): + params = litellm.get_optional_params( + model="gpt-6-astra", + custom_llm_provider="openai", + max_tokens=100, + reasoning_effort="max", + verbosity="low", + ) + assert params["max_completion_tokens"] == 100 + assert "max_tokens" not in params + assert params["reasoning_effort"] == "max" + assert params["verbosity"] == "low" diff --git a/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py b/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py index 1095819c98c..107a1afb2c6 100644 --- a/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py +++ b/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py @@ -41,6 +41,8 @@ from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config # Models that MUST be classified as GPT-5 (routed through GPT-5 reasoning path) GPT5_MODELS = [ + "gpt-6-astra", + "openai/gpt-6-astra", "gpt-5", "gpt-5.1", "gpt-5.2", @@ -120,6 +122,8 @@ class TestOpenAIGPT5ConfigIsModelGpt5Model: # /v1/responses bridge (when reasoning_effort is set and tools are passed) on # is_model_gpt_5_4_plus_model, so the gpt-5.6 family must land on the True side. GPT5_4_PLUS_MODELS = [ + "gpt-6-astra", + "openai/gpt-6-astra", "gpt-5.4", "gpt-5.5", "gpt-5.5-pro", diff --git a/tests/test_litellm/llms/openai/test_openai.py b/tests/test_litellm/llms/openai/test_openai.py new file mode 100644 index 00000000000..136b837f191 --- /dev/null +++ b/tests/test_litellm/llms/openai/test_openai.py @@ -0,0 +1,52 @@ +import pytest + +from litellm.llms.openai.openai import OpenAIChatCompletion + + +@pytest.mark.parametrize( + "api_base", + [ + None, + "https://api.openai.com/v1", + "https://api.openai.com:443/v1", + "https://southcentralus.privatelink.api.openai.com/v1", + "https://eu.api.openai.com/v1", + "https://us.api.openai.com/v1", + "HTTPS://API.OPENAI.COM/v1/", + ], +) +def test_get_stream_options_defaults_include_usage_on_every_openai_backed_host(api_base): + """ + PrivateLink and regional hostnames reach the real OpenAI backend, so a stream with no caller + stream_options must ask for the usage chunk exactly as the default base does. Regression guard + for LIT-6875: spend for those deployments fell back to local token counting. + """ + assert OpenAIChatCompletion().get_stream_options(stream_options=None, api_base=api_base) == { + "stream_options": {"include_usage": True} + } + + +@pytest.mark.parametrize( + "api_base", + [ + "https://my-gateway.example/v1", + "https://api.openai.com.evil.example/v1", + "https://notapi.openai.com/v1", + "https://gateway.example/v1?upstream=api.openai.com", + "https://openai.internal.example/api.openai.com/v1", + ], +) +def test_get_stream_options_leaves_foreign_hosts_without_a_usage_default(api_base): + """Only the host decides: an OpenAI-compatible backend elsewhere may not support stream_options at all.""" + assert OpenAIChatCompletion().get_stream_options(stream_options=None, api_base=api_base) == {} + + +@pytest.mark.parametrize( + "api_base", + ["https://southcentralus.privatelink.api.openai.com/v1", "https://my-gateway.example/v1"], +) +def test_get_stream_options_passes_caller_stream_options_through_on_any_host(api_base): + caller_options = {"include_usage": False} + assert OpenAIChatCompletion().get_stream_options(stream_options=caller_options, api_base=api_base) == { + "stream_options": caller_options + } diff --git a/tests/test_litellm/llms/openai/test_openai_common_utils.py b/tests/test_litellm/llms/openai/test_openai_common_utils.py index 3ae29e411e8..d3c21c5bd5a 100644 --- a/tests/test_litellm/llms/openai/test_openai_common_utils.py +++ b/tests/test_litellm/llms/openai/test_openai_common_utils.py @@ -7,7 +7,7 @@ import pytest import litellm from litellm.litellm_core_utils.token_counter import token_counter -from litellm.llms.openai.common_utils import BaseOpenAILLM +from litellm.llms.openai.common_utils import BaseOpenAILLM, is_openai_backed_api_base # Test parameters for different API functions API_FUNCTION_PARAMS = [ @@ -392,3 +392,22 @@ async def test_async_genuine_bad_request_still_raises(provider, stream): with pytest.raises(litellm.BadRequestError): await _call_and_drain() + + +@pytest.mark.parametrize( + ("api_base", "expected"), + [ + ("https://api.openai.com/v1", True), + ("https://api.openai.com:443/v1/", True), + ("https://southcentralus.privatelink.api.openai.com/v1", True), + ("https://eu.api.openai.com/v1", True), + ("HTTPS://API.OPENAI.COM/v1", True), + ("https://my-gateway.example/v1", False), + ("https://api.openai.com.evil.example/v1", False), + ("https://notapi.openai.com/v1", False), + ("https://gateway.example/v1?upstream=api.openai.com", False), + ("not a url", False), + ], +) +def test_is_openai_backed_api_base_decides_by_hostname_only(api_base, expected): + assert is_openai_backed_api_base(api_base) is expected diff --git a/tests/test_litellm/proxy/auth/test_litellm_license.py b/tests/test_litellm/proxy/auth/test_litellm_license.py index 8da365cb587..1db53638070 100644 --- a/tests/test_litellm/proxy/auth/test_litellm_license.py +++ b/tests/test_litellm/proxy/auth/test_litellm_license.py @@ -2,6 +2,8 @@ import asyncio import json from unittest.mock import AsyncMock, MagicMock, patch +from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey + from litellm.proxy.auth.litellm_license import LicenseCheck @@ -30,3 +32,70 @@ def test_is_over_limit(): assert license_check.is_over_limit(101) is False assert license_check.is_over_limit(100) is False assert license_check.is_over_limit(99) is False + + +def test_heuristic_v2_router_limit() -> None: + """Only the signed license's auto_router feature lifts the one-router limit; an API-verified + license (no airgapped data) and an airgapped license without the feature keep it.""" + license_check = LicenseCheck() + license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["auto_router"]} + assert license_check.heuristic_v2_router_limit() is None + + license_check.airgapped_license_data = { + "expiration_date": "2999-01-01", + "allowed_features": ["sso", "auto_router", "audit_logs"], + } + assert license_check.heuristic_v2_router_limit() is None + + license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["sso"]} + assert license_check.heuristic_v2_router_limit() == 1 + + license_check.airgapped_license_data = {"expiration_date": "2999-01-01"} + assert license_check.heuristic_v2_router_limit() == 1 + + license_check.airgapped_license_data = None + assert license_check.heuristic_v2_router_limit() == 1 + + +def _signed_license(expiration_date: str) -> tuple[RSAPublicKey, str]: + import base64 + + from cryptography.hazmat.primitives import hashes + from cryptography.hazmat.primitives.asymmetric import padding, rsa + + private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + message = json.dumps( + {"expiration_date": expiration_date, "user_id": "u", "allowed_features": ["auto_router"]} + ).encode() + signature = private_key.sign( + message, + padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH), + hashes.SHA256(), + ) + return private_key.public_key(), base64.b64encode(message + b"." + signature).decode() + + +def test_expired_or_unreadable_license_grants_no_features() -> None: + """The verifier stores the signed payload only after the expiry check passes and clears it when a + later verify rejects the license, so a stale payload cannot keep lifting the heuristic_v2 limit.""" + license_check = LicenseCheck() + public_key, valid_key = _signed_license("2999-01-01") + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True + assert license_check.heuristic_v2_router_limit() is None + + _, expired_key = _signed_license("2000-01-01") + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=expired_key) is not True + assert license_check.airgapped_license_data is None + assert license_check.heuristic_v2_router_limit() == 1 + + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True + assert license_check.verify_license_without_api_request(public_key=public_key, license_key="not-a-license") is not True + assert license_check.airgapped_license_data is None + + +def test_valid_signed_license_with_auto_router_lifts_the_limit() -> None: + license_check = LicenseCheck() + public_key, license_key = _signed_license("2999-01-01") + + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=license_key) is True + assert license_check.heuristic_v2_router_limit() is None diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index c04fb7b30ec..400eaf8ab3f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -193,6 +193,33 @@ async def test_apply_guardrail_compresses_and_returns_structured_messages( assert "headroom" in _applied_guardrails(request_data) +@pytest.mark.asyncio +async def test_apply_guardrail_leaves_background_requests_uncompressed( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + request_data = {"model": "gpt-4o", "background": True} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=_make_compress_response(COMPRESSED_MESSAGES), + ) as post: + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert result is inputs + post.assert_not_awaited() + assert _recorded_guardrail_entries(request_data) == [] + + def _recorded_guardrail_response(request_data: dict) -> dict: entries = request_data["metadata"]["standard_logging_guardrail_information"] assert len(entries) == 1 @@ -954,6 +981,40 @@ async def test_passthrough_handler_does_not_log_headroom_as_run( assert "headroom" not in _applied_guardrails(data) +@pytest.mark.asyncio +async def test_responses_request_sends_compressed_input_and_retrieve_tool_upstream( + guardrail: HeadroomGuardrail, +): + """Regression for LIT-6494: on /v1/responses the compressed messages must be + written back into `input`, not only the retrieve tool into `tools`, or the + model keeps reading the full document and never calls headroom_retrieve.""" + from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler + + data = { + "model": "gpt-5.6", + "instructions": ORIGINAL_MESSAGES[0]["content"], + "input": [{"role": m["role"], "content": m["content"]} for m in ORIGINAL_MESSAGES[1:]], + "tools": [{"type": "function", "name": "get_weather", "parameters": {"type": "object", "properties": {}}}], + } + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH), + ): + result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert result["instructions"] == ORIGINAL_MESSAGES[0]["content"] + assert [item["content"] for item in result["input"]] == [ + COMPRESSED_MESSAGES_WITH_HASH[0]["content"], + ORIGINAL_MESSAGES[2]["content"], + ORIGINAL_MESSAGES[3]["content"], + ] + assert "A" * 5000 not in json.dumps(result["input"]) + assert [tool["name"] for tool in result["tools"]] == ["get_weather", HEADROOM_RETRIEVE_TOOL_NAME] + + @pytest.mark.asyncio async def test_apply_guardrail_http_error_raises(): guardrail = _make_guardrail() @@ -1950,6 +2011,58 @@ def _openai_text_payload(content: str) -> dict: return _openai_completion_payload({"role": "assistant", "content": content}, "stop") +def _responses_retrieve_tool_definition() -> dict: + return {"type": "function", **_retrieve_tool_definition()["function"]} + + +def _openai_responses_payload(output_item: dict) -> dict: + return { + "id": "resp_ccr", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-4o", + "output": [output_item], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": {}, + "temperature": 1.0, + "top_p": 1.0, + "text": {"format": {"type": "text"}}, + "truncation": "disabled", + } + + +def _openai_responses_retrieve_call_payload() -> dict: + return _openai_responses_payload( + { + "type": "function_call", + "id": "fc_ccr", + "call_id": "call_ccr", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": json.dumps({"hash": CCR_HASH}), + "status": "completed", + } + ) + + +def _openai_responses_text_payload(text: str) -> dict: + return _openai_responses_payload( + { + "type": "message", + "id": "msg_ccr", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ) + + @pytest.mark.parametrize( "call_type, stream, tools, expect_conversion", [ @@ -1958,12 +2071,14 @@ def _openai_text_payload(content: str) -> dict: (CallTypes.acompletion, False, [_retrieve_tool_definition()], False), (CallTypes.acompletion, True, [{"type": "function", "function": {"name": "get_weather"}}], False), (CallTypes.acompletion, True, None, False), - (CallTypes.aresponses, True, [_retrieve_tool_definition()], False), + (CallTypes.aresponses, True, [_retrieve_tool_definition()], True), + (CallTypes.responses, True, [_responses_retrieve_tool_definition()], True), + (CallTypes.aresponses, False, [_retrieve_tool_definition()], False), (CallTypes.anthropic_messages, True, [_retrieve_tool_definition()], False), ], ) @pytest.mark.asyncio -async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions( +async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions_and_responses( guardrail: HeadroomGuardrail, call_type: CallTypes, stream: bool, @@ -1986,6 +2101,22 @@ async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_comple assert kwargs["stream"] is True +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_leaves_background_streams_alone(guardrail: HeadroomGuardrail): + kwargs = { + "model": "gpt-4o", + "stream": True, + "background": True, + "tools": [_responses_retrieve_tool_definition()], + } + + result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.aresponses) + + assert result is kwargs + assert HEADROOM_CONVERTED_STREAM_KEY not in kwargs + assert kwargs["stream"] is True + + @pytest.mark.asyncio async def test_pre_call_deployment_hook_still_compresses_for_deployment_level_configs( guardrail: HeadroomGuardrail, @@ -2094,6 +2225,117 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( assert not any(key.startswith("_headroom_interception") for key in followup_body) +@pytest.mark.asyncio +async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + """Regression test for LIT-6481: streaming /v1/responses must resolve the + retrieve tool call server-side exactly like streaming /chat/completions does, + instead of streaming a headroom_retrieve function_call to the client.""" + original_content = "the full uncompressed document" + final_answer = "the document says hello" + guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( + frozenset({CCR_HASH}), + time.monotonic() + 999, + ) + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[ + httpx.Response(200, json=_openai_responses_retrieve_call_payload()), + httpx.Response(200, json=_openai_responses_text_payload(final_answer)), + ] + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response(original_content), + ) as mock_get: + response = await litellm.aresponses( + model="openai/gpt-4o", + input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], + tools=[_responses_retrieve_tool_definition()], + stream=True, + litellm_call_id="ccr-call-id", + ) + events = [event async for event in response] + + streamed_text = "".join( + getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta" + ) + assert streamed_text == final_answer + assert not any("function_call" in str(getattr(event, "type", "")) for event in events) + assert not any( + getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events + ) + mock_get.assert_called_once() + assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + + assert len(upstream.calls) == 2 + followup_body = json.loads(upstream.calls[1].request.content) + assert not followup_body.get("stream") + assert original_content in json.dumps(followup_body["input"]) + assert not any(key.startswith("_headroom_interception") for key in followup_body) + + +def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + """The synchronous responses() path converts the stream the same way, so it + must hand back a stream iterator with the resolved answer rather than the + completed response object.""" + original_content = "the full uncompressed document" + final_answer = "the document says hello" + guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( + frozenset({CCR_HASH}), + time.monotonic() + 999, + ) + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[ + httpx.Response(200, json=_openai_responses_retrieve_call_payload()), + httpx.Response(200, json=_openai_responses_text_payload(final_answer)), + ] + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response(original_content), + ) as mock_get: + response = litellm.responses( + model="openai/gpt-4o", + input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], + tools=[_responses_retrieve_tool_definition()], + stream=True, + litellm_call_id="ccr-call-id", + ) + events = list(response) + + streamed_text = "".join( + getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta" + ) + assert streamed_text == final_answer + assert not any( + getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events + ) + mock_get.assert_called_once() + assert len(upstream.calls) == 2 + assert not json.loads(upstream.calls[1].request.content).get("stream") + + # --------------------------------------------------------------------------- # LIT-5018: the turn the model is being asked to act on is never compressed. # diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 68704f476b5..7954a4693cc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -17505,6 +17505,18 @@ def test_generate_key_request_blank_team_id_is_personal(): assert GenerateKeyRequest(team_id="team-1").team_id == "team-1" +def test_generate_key_request_blank_organization_and_project_id_are_unset(): + from litellm.proxy._types import RegenerateKeyRequest + + cleared = GenerateKeyRequest(organization_id="", project_id="") + assert cleared.organization_id is None + assert cleared.project_id is None + assert "organization_id" not in cleared.model_dump(exclude_none=True) + assert RegenerateKeyRequest(organization_id="").organization_id is None + assert GenerateKeyRequest(organization_id="org-1", project_id="proj-1").organization_id == "org-1" + assert GenerateKeyRequest(organization_id="org-1", project_id="proj-1").project_id == "proj-1" + + def test_key_request_blank_organization_id_is_unset(): from litellm.proxy._types import RegenerateKeyRequest, UpdateKeyRequest diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 5fa59a85c9d..c69f8f20a13 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -28,9 +28,18 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_team_models, ) from litellm.proxy.utils import PrismaClient +from litellm.router import Router from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment +async def _passthrough_row(update_data): + return update_data + + +async def _write_empty_row(**kwargs): + return await kwargs["write_row"]({}) + + class MockPrismaClient: def __init__( self, @@ -1191,7 +1200,7 @@ class TestTeamModelSiblingRouting: team_id = "team_no_alias" public_name = "gpt-4.1-mini" - async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client): + async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client, slot=None): return MagicMock(model_id=str(uuid.uuid4())) mock_team_model_add = AsyncMock() @@ -1372,7 +1381,8 @@ class TestTeamModelUpdate: db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, # type: ignore + prisma_client=prisma_client, # type: ignore, + write_row=_passthrough_row, ) assert result.get("model_name", "").startswith("model_name_test_team_123_") @@ -1435,7 +1445,6 @@ class TestTeamModelUpdate: team_id="team_123", public_model_name="new-public-name", db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, # type: ignore ) @@ -1481,7 +1490,6 @@ class TestTeamModelUpdate: team_id="team_123", public_model_name="new-public-name", db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=None, ) @@ -1490,39 +1498,72 @@ class TestTeamModelUpdate: mock_delete.assert_not_called() @pytest.mark.asyncio - async def test_rename_with_prisma_none_clears_patch_model_name(self): - """Rename path must clear patch_data.model_name even when prisma is unavailable (P1).""" + async def test_a_refused_row_write_leaves_the_team_untouched(self): + """The team's model list autocommits, so it is written only after the row write succeeded: a + refused write (the heuristic_v2 slot 403, a DB error) must not leave the team listing a name + whose row never changed.""" + from fastapi import HTTPException + from litellm.proxy.management_endpoints.model_management_endpoints import ( - _update_existing_team_model_assignment, + _update_team_model_in_db, ) from litellm.types.router import ModelInfo db_model = Deployment( - model_name="model_name_team_123_uuid1", + model_name="gpt-4o", litellm_params=LiteLLM_Params(model="azure/gpt-4o-mini"), - model_info=ModelInfo( - team_id="team_123", team_public_model_name="old-public-name" + model_info=ModelInfo(), + ) + user_api_key_dict = UserAPIKeyAuth(user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN) + events: list[str] = [] + written: dict[str, object] = {} + + def patch_data() -> updateDeployment: + return updateDeployment(model_name="team-public", model_info=ModelInfo(team_id="team_123")) + + async def refuse_row(update_data): + events.append("row") + raise HTTPException(status_code=403, detail="slot held") + + async def accept_row(update_data): + events.append("row") + written.update(update_data) + return update_data + + async def team_add(**_): + events.append("team_model_add") + + with ( + patch( # test-quality-ok: the team auth check needs a live DB; the write order is what is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.allow_team_model_action", + AsyncMock(return_value=True), ), - ) - patch_data = updateDeployment( - model_name="new-public-name", - model_info=ModelInfo(team_id="team_123"), - ) - user_api_key_dict = UserAPIKeyAuth( - user_id="test_user", - user_role=LitellmUserRoles.PROXY_ADMIN, - ) + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: team models are premium-gated through a proxy global with no injection seam + patch( # test-quality-ok: the team list write is the collaborator whose ordering is asserted + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add", + side_effect=team_add, + ), + ): + with pytest.raises(HTTPException): + await _update_team_model_in_db( + db_model=db_model, + patch_data=patch_data(), + user_api_key_dict=user_api_key_dict, + prisma_client=MockPrismaClient(team_exists=True), # type: ignore + write_row=refuse_row, + ) + assert events == ["row"] - await _update_existing_team_model_assignment( - team_id="team_123", - public_model_name="new-public-name", - db_model=db_model, - patch_data=patch_data, - user_api_key_dict=user_api_key_dict, - prisma_client=None, - ) - - assert patch_data.model_name is None + await _update_team_model_in_db( + db_model=db_model, + patch_data=patch_data(), + user_api_key_dict=user_api_key_dict, + prisma_client=MockPrismaClient(team_exists=True), # type: ignore + write_row=accept_row, + ) + assert events == ["row", "row", "team_model_add"] + assert str(written["model_name"]).startswith("model_name_team_123_") + assert "team-public" in str(written["model_info"]) @pytest.mark.asyncio async def test_rename_handles_legacy_string_model_info(self): @@ -1574,7 +1615,6 @@ class TestTeamModelUpdate: team_id="team_123", public_model_name="new-public-name", db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, # type: ignore ) @@ -1614,7 +1654,8 @@ class TestTeamModelUpdate: db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, # type: ignore + prisma_client=prisma_client, # type: ignore, + write_row=_passthrough_row, ) assert "403" in str(exc_info.value) @@ -1900,7 +1941,8 @@ class TestTeamModelUpdate: db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, # type: ignore + prisma_client=prisma_client, # type: ignore, + write_row=_passthrough_row, ) # team ACL must not be touched on a no-op edit @@ -4311,6 +4353,321 @@ class TestStrategyRouterWriteValidation: is None ) + @staticmethod + def _live_router_holding_one_heuristic_v2(limit: int | None) -> Router: + return Router( + model_list=[ + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}}, + { + "model_name": "held-v2", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}, + }, + "model_info": {"id": "held-id"}, + }, + ], + heuristic_v2_router_limit=lambda: limit, + ) + + class _FakeTx: + """Stands in for a prisma transaction: records the raw statements and exposes the model table.""" + + def __init__(self, db_held: int) -> None: + self.db_held = db_held + self.raw_calls: list[tuple[str, tuple[object, ...]]] = [] + self.litellm_proxymodeltable = MagicMock(create=AsyncMock(), update=AsyncMock()) + + async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]: + self.raw_calls.append((sql, args)) + return [{"held": self.db_held}] if "count(*)" in sql else [] + + async def __aenter__(self) -> "TestStrategyRouterWriteValidation._FakeTx": + return self + + async def __aexit__(self, *exc: object) -> None: + return None + + class _FakeDb: + """Stands in for prisma_client: the plain client and the transaction it opens are told apart by identity.""" + + def __init__(self, db_held: int, existing_row: object = None) -> None: + self.db = self + self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_held) + self.litellm_proxymodeltable = MagicMock( + create=AsyncMock(), update=AsyncMock(), find_unique=AsyncMock(return_value=existing_row) + ) + + def tx(self) -> "TestStrategyRouterWriteValidation._FakeTx": + return self.tx_obj + + _V2 = {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}} + _V1 = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini"}} + + @pytest.mark.parametrize( + "incoming,existing,expected", + [ + (_V2, None, _V2), + (_V2, _V1, _V2), + (None, _V1, _V1), + (None, None, None), + ("no-config", _V2, _V2), + ], + ) + def test_effective_complexity_router_config( + self, incoming: object, existing: object, expected: object + ) -> None: + """A write is judged on the config it leaves on the row: the incoming one when it carries one, else the stored one.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _effective_complexity_router_config, + ) + from litellm.types.router import updateLiteLLMParams + + incoming_params = None if incoming is None else updateLiteLLMParams( + complexity_router_config=None if incoming == "no-config" else incoming + ) + existing_params = None if existing is None else updateLiteLLMParams(complexity_router_config=existing) + assert _effective_complexity_router_config(incoming_params, existing_params) == expected + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "limit,effective_config,db_held,config_holds_one,model_id,expected", + [ + (1, _V2, 1, False, None, "refused"), + (1, _V2, 0, True, None, "refused"), + (1, _V2, 0, False, None, "reserved"), + (1, _V2, 0, False, "held-id", "reserved"), + (2, _V2, 1, False, None, "reserved"), + (1, _V1, 5, True, None, "plain"), + (1, None, 5, True, None, "plain"), + (None, _V2, 5, True, None, "plain"), + ], + ) + async def test_heuristic_v2_slot_matrix( + self, + limit: int | None, + effective_config: object, + db_held: int, + config_holds_one: bool, + model_id: str | None, + expected: str, + ) -> None: + """The slot is claimed inside a locked transaction only for a heuristic_v2 write under a limit; the DB rows + (other pods included) plus config.yaml routers decide, the row being edited is excluded through the SQL + parameter, and every other write runs on the plain client with no lock.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_management_endpoints import ( + HEURISTIC_V2_SLOT_LOCK_KEY, + _heuristic_v2_slot, + ) + + fake = self._FakeDb(db_held) + live_router = self._live_router_holding_one_heuristic_v2(limit) if config_holds_one else None + with ( + patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch("litellm.proxy.proxy_server.llm_router", live_router), # test-quality-ok: the guard reads the proxy router global with no injection seam + patch( # test-quality-ok: the cross-pod publish is the side effect under test; redis is not configured here + "litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change", + new=AsyncMock(), + ) as published, + ): + if expected == "refused": + with pytest.raises(HTTPException) as exc_info: + async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id): + pass + assert exc_info.value.status_code == 403 + assert "At most 1 auto-router" in str(exc_info.value.detail) + assert "'auto_router' feature lifts the limit" in str(exc_info.value.detail) + return + async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id) as tables: + handle = tables + if expected == "plain": + await handle.create(data={}) + fake.litellm_proxymodeltable.create.assert_awaited_once_with(data={}) + assert fake.tx_obj.raw_calls == [] + return + assert handle is fake.tx_obj.litellm_proxymodeltable + published.assert_awaited_once_with(redis_cache=None, object_type="litellm_proxymodeltable") + (lock_sql, lock_params), (_count_sql, count_params) = fake.tx_obj.raw_calls + assert "pg_advisory_xact_lock($1)" in lock_sql and "count" not in lock_sql + assert lock_params == (HEURISTIC_V2_SLOT_LOCK_KEY,) + assert count_params == (model_id or "",) + + @pytest.mark.asyncio + async def test_team_model_bookkeeping_runs_after_the_slot_is_released(self) -> None: + """team_model_add needs a second pool connection, so it must run only after the slot transaction + (and its advisory lock) has closed; a pool-sized burst of team creates would otherwise stall on the + lock holder waiting for a connection the waiters are occupying.""" + from contextlib import asynccontextmanager + + from litellm.proxy.management_endpoints.model_management_endpoints import _add_team_model_to_db + from litellm.types.router import ModelInfo + + events: list[str] = [] + created = MagicMock(model_id="row-1") + + @asynccontextmanager + async def slot(): + events.append("slot-enter") + yield MagicMock(create=AsyncMock(return_value=created)) + events.append("slot-exit") + + async def team_model_add(**_: object) -> None: + events.append("team_model_add") + + deployment = Deployment( + model_name="public-v2", + litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=self._V2), + model_info=ModelInfo(id="row-1", team_id="team-1"), + ) + with ( + patch( # test-quality-ok: params are encrypted with the proxy master key, which this test does not configure + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + lambda value, new_encryption_key=None: value, + ), + patch( # test-quality-ok: the team list write is the collaborator whose ordering is asserted + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add", + side_effect=team_model_add, + ), + ): + result = await _add_team_model_to_db( + model_params=deployment, + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + prisma_client=MagicMock(), + slot=slot(), + ) + + assert result is created + assert events == ["slot-enter", "slot-exit", "team_model_add"] + + @pytest.mark.asyncio + async def test_add_new_model_refuses_a_second_heuristic_v2_router_before_the_db_write(self) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + add_new_model, + ) + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + fake = self._FakeDb(db_held=1) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( # test-quality-ok: params are encrypted before the slot is entered; no master key in this test + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + lambda value, new_encryption_key=None: value, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model( + model_params=Deployment( + model_name="second-v2", + litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=self._V2), + ), + user_api_key_dict=admin, + ) + assert exc_info.value.code == "403" + assert "At most 1 auto-router" in str(exc_info.value.message) + fake.tx_obj.litellm_proxymodeltable.create.assert_not_awaited() + fake.litellm_proxymodeltable.create.assert_not_awaited() + + @pytest.mark.asyncio + async def test_patch_model_refuses_switching_another_router_to_heuristic_v2(self) -> None: + """patch_model relays HTTPException as-is, so the license refusal reaches the client as a plain 403.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_management_endpoints import ( + patch_model, + ) + from litellm.types.router import updateLiteLLMParams + + model_id = "other-id" + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + fake = self._FakeDb(db_held=1) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( # test-quality-ok: the write must be refused before this DB step runs + "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", + new=AsyncMock(return_value=self._db_complexity_router(model_id)), + ), + patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( # test-quality-ok: the helper's team bookkeeping needs a live DB; the row writer it is handed is what is under test + "litellm.proxy.management_endpoints.model_management_endpoints._update_team_model_in_db", + new=AsyncMock(side_effect=_write_empty_row), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await patch_model( + model_id=model_id, + patch_data=updateDeployment(litellm_params=updateLiteLLMParams(complexity_router_config=self._V2)), + user_api_key_dict=admin, + ) + assert exc_info.value.status_code == 403 + fake.tx_obj.litellm_proxymodeltable.update.assert_not_awaited() + fake.litellm_proxymodeltable.update.assert_not_awaited() + + @pytest.mark.asyncio + async def test_update_model_refuses_switching_another_router_to_heuristic_v2(self) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_model, + ) + from litellm.types.router import ModelInfo, updateLiteLLMParams + + model_id = "other-id" + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + existing_row = MagicMock() + existing_row.model_dump.return_value = { + "model_name": "my-auto-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}}, + }, + "model_info": {"id": model_id}, + } + existing_row.litellm_params = existing_row.model_dump.return_value["litellm_params"] + fake = self._FakeDb(db_held=1, existing_row=existing_row) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + ): + with pytest.raises(ProxyException) as exc_info: + await update_model( + model_params=updateDeployment( + litellm_params=updateLiteLLMParams(complexity_router_config=self._V2), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=admin, + ) + assert exc_info.value.code == "403" + fake.tx_obj.litellm_proxymodeltable.update.assert_not_awaited() + fake.litellm_proxymodeltable.update.assert_not_awaited() + @pytest.mark.asyncio async def test_update_model_rejects_prefix_strip(self): from litellm.proxy._types import ProxyException diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index a1c38d26b9d..d35b77f732c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -45,6 +45,10 @@ from litellm.types.router import ( from litellm.types.utils import Usage +async def _passthrough_row(update_data): + return update_data + + def test_model_info_accepts_valid_ptu_fields(): info = ModelInfo( id="x", @@ -385,6 +389,7 @@ class TestTeamModelUpdateValidatesBeforeWriting: patch_data=patch_data, user_api_key_dict=MagicMock(), prisma_client=MagicMock(), + write_row=_passthrough_row, ) return result, touched @@ -914,6 +919,7 @@ class TestPtuDeploymentsAreNotBilledPerToken: patch_data=patch, user_api_key_dict=UserAPIKeyAuth(user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN), prisma_client=MagicMock(), + write_row=_passthrough_row, ) assert exc.value.status_code == 400 diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 088e82b370f..166fcb2863c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -9897,11 +9897,11 @@ class TestResolveTeamAccessGroupResources: assert resolved.access_group_mcp_server_ids == ["mcp-1"] assert resolved.access_group_agent_ids == ["agent-1"] assert [ - (d.access_group_id, d.access_group_name, d.models) + (d.access_group_id, d.access_group_name, d.models, d.mcp_server_ids, d.agent_ids) for d in (resolved.access_group_details or []) ] == [ - ("ag-1", "shared-models", ("gpt-4", "claude-3")), - ("ag-2", "extra-models", ("claude-3", "gemini")), + ("ag-1", "shared-models", ("gpt-4", "claude-3"), ("mcp-1",), ()), + ("ag-2", "extra-models", ("claude-3", "gemini"), (), ("agent-1",)), ] @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py index 961479c0393..6735f2a3780 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py @@ -5,6 +5,7 @@ import pytest from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _base_vertex_proxy_route, + _upstream_headers_for_vertex_route, ) from litellm.types.router import DeploymentTypedDict @@ -348,6 +349,93 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header(): assert headers_passed_through is False +VERTEX_ANTHROPIC_MODELS_PREFIX = "v1/projects/test-project/locations/global/publishers/anthropic/models/" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_segment", "expects_anthropic_beta"), + [ + ("count-tokens:rawPredict", False), + ("claude-sonnet-4-6:streamRawPredict", True), + ], +) +async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens( + model_segment: str, expects_anthropic_beta: bool +): + with ( + patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it + "litellm.proxy.proxy_server.llm_router", None + ), + patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router" + ) as mock_pt_router, + patch( # test-quality-ok: the route offers no injection point for its header preparation + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers", + new_callable=AsyncMock, + ) as mock_prep_headers, + patch( # test-quality-ok: the upstream call is captured here, the route offers no injection point + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" + ) as mock_create_route, + patch( # test-quality-ok: the route calls auth directly rather than through Depends + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + patch( # test-quality-ok: the route reads the request body for this, a MagicMock request has none + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn", + new_callable=AsyncMock, + return_value=False, + ), + ): + mock_pt_router.get_vertex_credentials.return_value = MagicMock() + mock_prep_headers.return_value = ( + { + "anthropic-beta": "tool-search-tool-2025-10-19,web-search-2025-03-05", + "content-type": "application/json", + "Authorization": "Bearer vertex-access-token", + }, + "https://aiplatform.googleapis.com", + False, + "test-project", + "global", + ) + mock_create_route.return_value = AsyncMock() + mock_auth.return_value = UserAPIKeyAuth(api_key="sk-litellm-secret-key") + + await _base_vertex_proxy_route( + endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}", + request=MagicMock(), + fastapi_response=MagicMock(), + get_vertex_pass_through_handler=MagicMock(), + ) + + upstream_headers = mock_create_route.call_args.kwargs["custom_headers"] + assert ("anthropic-beta" in upstream_headers) is expects_anthropic_beta + assert upstream_headers["Authorization"] == "Bearer vertex-access-token" + assert upstream_headers["content-type"] == "application/json" + + +def test_upstream_headers_for_vertex_route_filters_anthropic_beta_by_route(): + headers = { + "Anthropic-Beta": "effort-2025-11-24", + "content-type": "application/json", + "Authorization": "Bearer vertex-access-token", + } + + count_tokens_headers = _upstream_headers_for_vertex_route( + f"{VERTEX_ANTHROPIC_MODELS_PREFIX}count-tokens:rawPredict", headers + ) + model_headers = _upstream_headers_for_vertex_route( + f"{VERTEX_ANTHROPIC_MODELS_PREFIX}claude-sonnet-4-6:rawPredict", headers + ) + + assert dict(count_tokens_headers) == { + "content-type": "application/json", + "Authorization": "Bearer vertex-access-token", + } + assert dict(model_headers) == headers + + @pytest.mark.asyncio async def test_vertex_passthrough_does_not_forward_litellm_auth_token(): """ diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 1ab18639fff..dcfad8f6815 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -28,6 +28,7 @@ from litellm.proxy.proxy_server import ( resolve_routing_plugins, validate_deployment_complexity_router_placement, validate_deployment_max_agentic_loops, + validate_heuristic_v2_router_limit, ) from .conftest import normalize @@ -193,6 +194,120 @@ def test_validate_deployment_complexity_router_placement_leaves_valid_deployment assert model["litellm_params"] == litellm_params +def _heuristic_v2_row(model_name: str, classifier_type: str = "heuristic_v2") -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"classifier_type": classifier_type, "tiers": {"SIMPLE": "gpt-4o-mini"}}, + }, + } + + +def test_validate_heuristic_v2_router_limit_refuses_to_start_over_the_limit() -> None: + """Same reason as the two validators above: the proxy router swallows registration errors, so + an over-limit config.yaml must fail here instead of booting with a silently missing router.""" + with pytest.raises(ValueError, match=re.escape("At most 1 auto-router")) as exc_info: + validate_heuristic_v2_router_limit( + [_heuristic_v2_row("a"), _heuristic_v2_row("b"), _heuristic_v2_row("c", "heuristic")], limit=1 + ) + assert "'auto_router' feature lifts the limit" in str(exc_info.value) + + +@pytest.mark.parametrize( + "model_list,limit", + [ + ([_heuristic_v2_row("a"), _heuristic_v2_row("b")], None), + ([_heuristic_v2_row("a"), _heuristic_v2_row("c", "heuristic")], 1), + ([{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}], 1), + ], +) +def test_validate_heuristic_v2_router_limit_leaves_configs_within_the_limit_alone( + model_list: list[dict[str, object]], limit: int | None +) -> None: + assert validate_heuristic_v2_router_limit(model_list, limit=limit) is None + + +_TWO_HEURISTIC_V2_ROUTERS_YAML = ( + "model_list:\n" + " - model_name: gpt-4o-mini\n" + " litellm_params:\n" + " model: openai/gpt-4o-mini\n" + " api_key: k\n" + " - model_name: v2-a\n" + " litellm_params:\n" + " model: auto_router/complexity_router\n" + " complexity_router_config:\n" + " classifier_type: heuristic_v2\n" + " tiers: {SIMPLE: gpt-4o-mini}\n" + " - model_name: v2-b\n" + " litellm_params:\n" + " model: auto_router/complexity_router\n" + " complexity_router_config:\n" + " classifier_type: heuristic_v2\n" + " tiers: {SIMPLE: gpt-4o-mini}\n" + "router_settings:\n" + " heuristic_v2_router_limit: 99\n" +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("license_limit", [1, None]) +async def test_ProxyConfig_load_config_takes_the_heuristic_v2_limit_from_the_license_only( + tmp_path, monkeypatch, license_limit: int | None +) -> None: + """`router_settings.heuristic_v2_router_limit` is managed outside config.yaml: an operator + cannot grant the entitlement by editing the config, and a licensed proxy boots both routers.""" + f = tmp_path / "c.yaml" + f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setattr( + "litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: license_limit + ) + + if license_limit is None: + router, _model_list, _general_settings = await ProxyConfig().load_config( + router=None, config_file_path=str(f) + ) + assert router.heuristic_v2_router_limit is not None + assert router.heuristic_v2_router_limit() is None + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + return + + with pytest.raises(ValueError, match=re.escape("config.yaml model_list: At most 1 auto-router")): + await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + +@pytest.mark.asyncio +async def test_ProxyConfig_load_config_router_refuses_a_db_heuristic_v2_router_beyond_the_license( + tmp_path, monkeypatch +) -> None: + """config.yaml holds the one allowed heuristic_v2 router; a second one arriving later from the DB + is refused at registration because the router was built with the license's ceiling.""" + from litellm.types.router import Deployment + + f = tmp_path / "c.yaml" + f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML.replace(" - model_name: v2-b\n", " - model_name: v1-b\n", 1).replace( + "classifier_type: heuristic_v2\n tiers: {SIMPLE: gpt-4o-mini}\nrouter_settings", + "classifier_type: heuristic\n tiers: {SIMPLE: gpt-4o-mini}\nrouter_settings", + )) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setattr("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1) + + router, _model_list, _general_settings = await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + assert router.heuristic_v2_router_limit is not None + assert router.heuristic_v2_router_limit() == 1 + assert sorted(router.complexity_routers) == ["v1-b", "v2-a"] + db_row = Deployment(**_heuristic_v2_row("v2-from-db"), model_info={"id": "db-id"}) + assert router.upsert_deployment(db_row) is None + assert sorted(router.complexity_routers) == ["v1-b", "v2-a"] + + def test_validate_deployment_max_agentic_loops_allows_a_deployment_without_the_key(): model = {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}} diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 60a32102946..f4cd8814bc1 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -4234,6 +4234,91 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_spend(): assert call_args[2] == [api_key] +@pytest.mark.asyncio +async def test_build_ui_spend_logs_response_sums_multi_round_session_tokens(): + """ + Regression test for LIT-4929: the logs table showed the summed session cost but + only the last call's token usage. Every row of a multi-round session must carry + the session-wide prompt, completion and total token sums from the aggregate + query, while rows outside a session carry none of them. + """ + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + _build_ui_spend_logs_response, + ) + + session_id = "sess-multi-round-tokens" + api_key = "hashed-key-xyz" + dict_rows = [ + { + "request_id": "req-1", + "session_id": session_id, + "call_type": "completion", + "api_key": api_key, + "total_tokens": 10, + "prompt_tokens": 7, + "completion_tokens": 3, + }, + { + "request_id": "req-2", + "session_id": session_id, + "call_type": "completion", + "api_key": api_key, + "total_tokens": 50, + "prompt_tokens": 35, + "completion_tokens": 15, + }, + { + "request_id": "req-3", + "session_id": None, + "call_type": "completion", + "api_key": api_key, + "total_tokens": 5, + "prompt_tokens": 4, + "completion_tokens": 1, + }, + ] + + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + { + "session_id": session_id, + "api_key": api_key, + "session_total_count": 2, + "session_total_spend": 0.06, + "mcp_tool_call_count": 0, + "mcp_tool_call_spend": 0.0, + "session_total_prompt_tokens": 42, + "session_total_completion_tokens": 18, + "session_total_tokens": 60, + } + ] + ) + + result = await _build_ui_spend_logs_response( + prisma_client=mock_prisma, + data=dict_rows, + total_records=3, + page=1, + page_size=50, + total_pages=1, + enrich_session_counts=True, + ) + + rows = result["data"] + session_rows = rows[:2] + assert [row["session_total_tokens"] for row in session_rows] == [60, 60] + assert [row["session_total_prompt_tokens"] for row in session_rows] == [42, 42] + assert [row["session_total_completion_tokens"] for row in session_rows] == [18, 18] + assert [(row["total_tokens"], row["prompt_tokens"], row["completion_tokens"]) for row in session_rows] == [ + (10, 7, 3), + (50, 35, 15), + ] + + token_keys = ("session_total_tokens", "session_total_prompt_tokens", "session_total_completion_tokens") + assert all(key not in rows[2] for key in token_keys) + + @pytest.mark.asyncio async def test_build_ui_spend_logs_response_session_cache_hit_count(): """ diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index dcaad968663..f16c6c937d0 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -942,6 +942,47 @@ def test_create_model_info_response_uses_deployment_limits_when_not_in_cost_map( assert response["max_output_tokens"] == 8000 +def test_create_model_info_response_uses_deployment_mode_for_auto_router(): + router = litellm.Router( + model_list=[ + { + "model_name": "claude-sonnet", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, + }, + { + "model_name": "claude-auto", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": { + "SIMPLE": "claude-sonnet", + "MEDIUM": "claude-sonnet", + "COMPLEX": "claude-sonnet", + } + }, + "complexity_router_default_model": "claude-sonnet", + }, + "model_info": { + "mode": "chat", + "max_input_tokens": 1_000_000, + "max_output_tokens": 128_000, + }, + }, + ] + ) + + response = create_model_info_response( + model_id="claude-auto", + provider="openai", + llm_router=router, + get_model_info=_raise_unmapped, + ) + + assert response["mode"] == "chat" + assert response["max_input_tokens"] == 1_000_000 + assert response["max_output_tokens"] == 128_000 + + def test_create_model_info_response_deployment_limits_override_cost_map(): router = MagicMock() router.get_configured_token_limits.return_value = (200000, None) diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py index 7d72121456a..93049b21460 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py @@ -120,9 +120,7 @@ async def test_delete_vector_store_checks_access(): "team_id": "team_456", } ) - mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock( - return_value=mock_vector_store - ) + mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=mock_vector_store) # User from different team should get 403 user_api_key_dict = UserAPIKeyAuth(team_id="team_789") @@ -134,9 +132,115 @@ async def test_delete_vector_store_checks_access(): ): with patch("litellm.vector_store_registry", None): with pytest.raises(HTTPException) as exc_info: - await delete_vector_store( - data=request, user_api_key_dict=user_api_key_dict - ) + await delete_vector_store(data=request, user_api_key_dict=user_api_key_dict) assert exc_info.value.status_code == 403 assert "Access denied" in exc_info.value.detail + + +_UNSCOPED: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_unscoped", + "custom_llm_provider": "openai", + "team_id": None, +} +_TEAM_A_OWNED: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_team_a", + "custom_llm_provider": "openai", + "team_id": "team_a", +} +_UI_CREATED: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_ui_created", + "custom_llm_provider": "openai", + "team_id": "litellm-dashboard", +} + + +async def _listed_ids(user_api_key_dict: UserAPIKeyAuth) -> list[str]: + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + list_vector_stores, + ) + + with patch( # test-quality-ok: the list route reads rows through this module-level DB helper, no injection seam + "litellm.proxy.vector_store_endpoints.management_endpoints.VectorStoreRegistry._get_vector_stores_from_db", + new=AsyncMock(return_value=[_UNSCOPED, _TEAM_A_OWNED, _UI_CREATED]), + ): + response = await list_vector_stores(user_api_key_dict=user_api_key_dict) + return sorted(vs["vector_store_id"] for vs in response["data"]) + + +@pytest.mark.asyncio +async def test_list_vector_stores_hides_ungranted_stores_from_non_admin_keys(): + """A store with no team_id and no allowlist entry is not listed for a key it was never granted to; + only team ownership or an explicit object_permission grant makes a store visible.""" + assert await _listed_ids(UserAPIKeyAuth()) == [] + assert await _listed_ids(UserAPIKeyAuth(team_id="team_a")) == ["vs_team_a"] + assert await _listed_ids( + UserAPIKeyAuth( + team_id="team_b", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-1", vector_stores=["vs_unscoped"]), + ) + ) == ["vs_unscoped"] + assert await _listed_ids( + UserAPIKeyAuth( + team_id="team_b", + team_object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-2", vector_stores=["vs_unscoped"] + ), + ) + ) == ["vs_unscoped"] + assert await _listed_ids(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)) == [ + "vs_team_a", + "vs_ui_created", + "vs_unscoped", + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("user_team_ids", "session_key_grants", "expected"), + [ + ([], None, []), + ([], ["vs_unscoped"], ["vs_unscoped"]), + (["team_a"], None, ["vs_team_a"]), + (["team_a", "team_granted"], None, ["vs_team_a", "vs_unscoped"]), + ], +) +async def test_list_vector_stores_dashboard_session_resolves_real_teams( + user_team_ids: list[str], session_key_grants: list[str] | None, expected: list[str] +): + """A dashboard session lists through the user's real teams plus the session key's own grants: stores created + from the dashboard (team_id litellm-dashboard) are not visible just because every session shares that team id, + while stores owned by or granted to one of the user's teams, or granted to the session key itself, are.""" + from litellm.models.team import LiteLLM_TeamTableCachedObj + + alice = UserAPIKeyAuth( + team_id="litellm-dashboard", + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER, + object_permission=( + LiteLLM_ObjectPermissionTable(object_permission_id="op-4", vector_stores=session_key_grants) + if session_key_grants is not None + else None + ), + ) + teams = { + "team_a": LiteLLM_TeamTableCachedObj(team_id="team_a"), + "team_granted": LiteLLM_TeamTableCachedObj( + team_id="team_granted", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-3", vector_stores=["vs_unscoped"]), + ), + } + + async def fake_get_team_object(team_id: str, **_kwargs: object) -> LiteLLM_TeamTableCachedObj: + return teams[team_id] + + with ( + patch( # test-quality-ok: team rows come from the module-level prisma client, no injection seam + "litellm.proxy.auth.auth_checks.get_team_object", new=fake_get_team_object + ), + patch( # test-quality-ok: the user row comes from the module-level prisma client, no injection seam + "litellm.proxy.vector_store_endpoints.utils.resolve_ui_session_team_ids", + new=AsyncMock(return_value=user_team_ids), + ), + ): + assert await _listed_ids(alice) == expected diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index a3dd5688ad1..761e87ac764 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -1,6 +1,7 @@ import asyncio import time from types import TracebackType +from typing import Final from unittest.mock import MagicMock, patch @@ -294,3 +295,39 @@ async def test_azure_health_check_honors_deployment_realtime_protocol(): model_params={"realtime_protocol": "GA"}, ) assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-4o-realtime-preview" + + +class _ConnectThatStopsAfterCapturingTheUrl: + url: str | None = None + + def __call__(self, url: str, **kwargs: object) -> "_ConnectThatStopsAfterCapturingTheUrl": + self.url = url + return self + + async def __aenter__(self) -> None: + raise RuntimeError("backend url captured, nothing to bridge") + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + return None + + +@pytest.mark.asyncio +async def test_arealtime_azure_ai_on_a_foundry_host_connects_to_the_azure_openai_realtime_route(): + connect: Final = _ConnectThatStopsAfterCapturingTheUrl() + with patch("websockets.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="azure_ai/gpt-realtime-mini", + websocket=MagicMock(), + api_base="https://my-project.services.ai.azure.com", + api_key="fake-key", + litellm_logging_obj=FakeLogging(), + ) + assert connect.url == ( + "wss://my-project.services.ai.azure.com/openai/realtime" + "?api-version=2024-10-01-preview&deployment=gpt-realtime-mini" + ) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 4a03913f55a..719d51c11e3 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -25,6 +25,7 @@ from litellm.types.utils import ( ModelResponse, ModelResponseStream, StreamingChoices, + Usage, ) CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256" @@ -527,6 +528,20 @@ async def test_streaming_response_id_falls_back_when_upstream_yields_nothing(): assert response_ids[0].startswith("resp_") +def test_completed_event_restores_usage_hidden_by_stream_options_none(): + final_chunk = _chunk("", finish_reason="stop") + final_chunk._hidden_params = {"usage": Usage(prompt_tokens=117, completion_tokens=5, total_tokens=122)} + iterator = _build_iterator([_chunk("the document says hello"), final_chunk]) + + events = list(iterator) + + completed = next( + event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + assert completed.response.usage.input_tokens == 117 + assert completed.response.usage.output_tokens == 5 + + def test_object_tool_call_arguments_stream_as_valid_json(): """A provider that sends decoded object arguments must still stream valid JSON. diff --git a/tests/test_litellm/router_strategy/test_auto_router.py b/tests/test_litellm/router_strategy/test_auto_router.py index 36199b45847..123ada83ca4 100644 --- a/tests/test_litellm/router_strategy/test_auto_router.py +++ b/tests/test_litellm/router_strategy/test_auto_router.py @@ -330,36 +330,46 @@ ROUTER_CONFIG: Final = json.dumps( ) -class FailingRouteLayer: - """Route layer whose embedding call fails, as it does when the prompt exceeds the encoder's window.""" - - def __call__(self, text: str) -> Any: - raise ValueError( - "Internal_litellm_router API call failed. Error: litellm.InternalServerError: " - "input is too large to process. increase the physical batch size" - ) - - class FixedRouteLayer: - """Route layer that returns whatever the test tells it to, recording the text it was asked about.""" + """Route layer that returns whatever the test tells it to for the query vector it is handed.""" def __init__(self, route_choice: Any) -> None: self.route_choice = route_choice - self.seen_text: str | None = None - def __call__(self, text: str) -> Any: - self.seen_text = text + async def acall(self, vector: Any) -> Any: return self.route_choice +def _embedding_response(input: List[str]) -> Any: + import litellm + + return litellm.EmbeddingResponse( + data=[{"embedding": [0.1, 0.2], "index": i, "object": "embedding"} for i in range(len(input))] + ) + + class StubEmbeddingRouter: - """Stands in for the LiteLLM Router when the route index has to be built for real.""" + """Stands in for the LiteLLM Router, recording the text and kwargs each query embedding was made with.""" + + def __init__(self) -> None: + self.seen_text: str | None = None + self.aembedding_kwargs: Dict[str, Any] | None = None def embedding(self, input: List[str], model: str, **kwargs: Any) -> Any: - import litellm + return _embedding_response(input) - return litellm.EmbeddingResponse( - data=[{"embedding": [0.1, 0.2], "index": i, "object": "embedding"} for i in range(len(input))] + async def aembedding(self, input: List[str], model: str, **kwargs: Any) -> Any: + self.seen_text = input[0] + self.aembedding_kwargs = kwargs + return _embedding_response(input) + + +class FailingEmbeddingRouter(StubEmbeddingRouter): + """Router whose query embedding fails, as it does when the prompt exceeds the encoder's window.""" + + async def aembedding(self, input: List[str], model: str, **kwargs: Any) -> Any: + raise ValueError( + "litellm.InternalServerError: input is too large to process. increase the physical batch size" ) @@ -369,7 +379,7 @@ def _auto_router(routelayer: Any, litellm_router_instance: Any = None, **kwargs: auto_router_config=ROUTER_CONFIG, default_model="fallback-model", embedding_model="text-embedding-3-small", - litellm_router_instance=litellm_router_instance or MagicMock(), + litellm_router_instance=litellm_router_instance or StubEmbeddingRouter(), **kwargs, ) auto_router.routelayer = routelayer @@ -381,7 +391,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel: @pytest.mark.asyncio async def test_should_fall_back_to_default_model_when_the_embedding_call_fails(self): - auto_router: Final = _auto_router(FailingRouteLayer()) + auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=FailingEmbeddingRouter()) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -440,8 +450,8 @@ class TestAutoRouterAlwaysResolvesARoutableModel: async def test_should_still_route_to_the_matched_route_when_one_matches(self): from semantic_router.schema import RouteChoice - layer: Final = FixedRouteLayer(RouteChoice(name="code-model")) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -451,7 +461,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel: assert result is not None assert result.model == "code-model" - assert layer.seen_text == "fix this stack trace" + assert router.seen_text == "fix this stack trace" class TestAutoRouterEmbeddingInputCap: @@ -483,8 +493,8 @@ class TestAutoRouterRoutesResponsesApiInput: async def test_should_route_a_string_input_when_messages_is_none(self): from semantic_router.schema import RouteChoice - layer: Final = FixedRouteLayer(RouteChoice(name="code-model")) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -498,14 +508,14 @@ class TestAutoRouterRoutesResponsesApiInput: assert result is not None assert result.model == "code-model" assert result.messages is None - assert layer.seen_text == "fix this stack trace" + assert router.seen_text == "fix this stack trace" @pytest.mark.asyncio async def test_should_route_a_list_input_with_instructions_when_messages_is_none(self): from semantic_router.schema import RouteChoice - layer: Final = FixedRouteLayer(RouteChoice(name="code-model")) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -525,13 +535,13 @@ class TestAutoRouterRoutesResponsesApiInput: assert result is not None assert result.model == "code-model" - assert layer.seen_text is not None - assert "fix this stack trace" in layer.seen_text + assert router.seen_text is not None + assert "fix this stack trace" in router.seen_text @pytest.mark.asyncio async def test_should_skip_routing_when_neither_messages_nor_input_is_present(self): - layer: Final = FixedRouteLayer(None) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -540,12 +550,12 @@ class TestAutoRouterRoutesResponsesApiInput: ) assert result is None - assert layer.seen_text is None + assert router.seen_text is None @pytest.mark.asyncio async def test_should_keep_routing_an_empty_messages_list_to_the_default_model(self): - layer: Final = FixedRouteLayer(None) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -555,4 +565,42 @@ class TestAutoRouterRoutesResponsesApiInput: assert result is not None assert result.model == "fallback-model" - assert layer.seen_text == "" + assert router.seen_text == "" + + +class TestAutoRouterAttributesItsEmbeddingSpend: + """The query embedding is billed to the key that sent the request, like any other call it made.""" + + @pytest.mark.asyncio + async def test_should_forward_the_callers_identity_to_the_query_embedding_minus_its_budget_reservation(self): + from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY + + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(None, litellm_router_instance=router) + request_kwargs: Final = { + "metadata": { + "user_api_key": "hashed-key", + "user_api_key_team_id": "team-1", + "user_api_key_budget_reservation": {"reservation_id": "r-1"}, + }, + "litellm_session_id": "session-1", + } + + result: Final = await auto_router.async_pre_routing_hook( + model="my-auto-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "fix this stack trace"}], + ) + + assert result is not None + assert router.seen_text == "fix this stack trace" + assert router.aembedding_kwargs is not None + forwarded: Final = router.aembedding_kwargs["metadata"] + assert forwarded["user_api_key"] == "hashed-key" + assert forwarded["user_api_key_team_id"] == "team-1" + assert forwarded[INTERNAL_CALL_ORIGIN_METADATA_KEY] == "autorouter_classifier" + assert "user_api_key_budget_reservation" not in forwarded + assert router.aembedding_kwargs["litellm_session_id"] == "session-1" + assert router.aembedding_kwargs["proxy_server_request"] == { + "body": {"model": "text-embedding-3-small", "input": ["fix this stack trace"]} + } diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index aa1b51afe10..da3791da39a 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -14,6 +14,7 @@ from pydantic import ValidationError import litellm from litellm import Router +from litellm.router_utils.auto_router_model_naming import count_heuristic_v2_routers from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY @@ -1085,6 +1086,165 @@ class TestRouterComplexityDeploymentMethods: router.init_complexity_router_deployment(deployment) assert "auto_router/complexity_router/test-router" in router.complexity_routers + @staticmethod + def _router_row(model_name: str, model_id: str, classifier_type: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": classifier_type, + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + }, + }, + "model_info": {"id": model_id}, + } + + _POOL: dict[str, object] = { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}, + } + + def test_heuristic_v2_ceiling_keeps_the_first_router_and_drops_the_rest(self) -> None: + """The proxy runs with ignore_invalid_deployments, so the second heuristic_v2 router is dropped + at registration while a heuristic (v1) sibling and the first v2 router stay routable.""" + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + self._router_row("v1-c", "id-c", "heuristic"), + ], + heuristic_v2_router_limit=lambda: 1, + ignore_invalid_deployments=True, + ) + + assert sorted(router.complexity_routers) == ["v1-c", "v2-a"] + assert router.get_deployment(model_id="id-b") is None + + def test_heuristic_v2_ceiling_raises_without_ignore_invalid_deployments(self) -> None: + with pytest.raises(ValueError, match="At most 1 auto-router"): + Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: 1, + ) + + def test_heuristic_v2_limit_is_resolved_on_every_registration(self) -> None: + """The Router never caches the limit: when the resolver's answer moves (the proxy re-verified + its license), the next registration and the next limit query see the new value.""" + limits = {"value": None} + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: limits["value"], + ignore_invalid_deployments=True, + ) + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + assert router.heuristic_v2_router_limit_violation() is None + + limits["value"] = 1 + assert router.heuristic_v2_router_limit_violation() is not None + assert router.upsert_deployment(Deployment(**self._router_row("v2-c", "id-c", "heuristic_v2"))) is None + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + + def test_heuristic_v2_ceiling_tightening_refuses_the_edit_and_keeps_the_live_router(self) -> None: + """Two heuristic_v2 routers registered under an unlimited ceiling, then the ceiling drops to one: + an edit to either must be refused before its live row is popped, or the failed re-add and + the failed restore would drop a serving router while the write reports success.""" + limits = {"value": None} + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: limits["value"], + ignore_invalid_deployments=True, + ) + limits["value"] = 1 + + assert router.upsert_deployment(Deployment(**self._router_row("v2-a-renamed", "id-a", "heuristic_v2"))) is None + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + assert router.get_deployment(model_id="id-a") is not None + + assert router.upsert_deployment(Deployment(**self._router_row("v1-a", "id-a", "heuristic"))) is not None + assert sorted(router.complexity_routers) == ["v1-a", "v2-b"] + + def test_config_deployments_excludes_db_rows(self) -> None: + """The proxy counts config.yaml routers from here and DB rows from the database, so a DB-loaded + row (``model_info.db_model``) must not show up twice.""" + router = Router(model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")]) + db_row = self._router_row("v2-db", "id-db", "heuristic_v2") + db_row["model_info"] = {"id": "id-db", "db_model": True} + assert router.upsert_deployment(Deployment(**db_row)) is not None + + assert sorted(str(row["model_name"]) for row in router.config_deployments()) == ["gpt-4o-mini", "v2-a"] + assert count_heuristic_v2_routers(router.config_deployments()) == 1 + + def test_failed_edit_of_a_live_v2_router_rolls_back_without_the_ceiling(self) -> None: + """A rollback after a failed upsert re-admits state that was already serving, so it must not be + judged by a ceiling that tightened since: converting one of two live heuristic_v2 routers to a + config whose registration fails must leave it serving its previous v2 configuration.""" + limits = {"value": None} + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: limits["value"], + ignore_invalid_deployments=True, + ) + limits["value"] = 1 + + broken = self._router_row("v1-a", "id-a", "heuristic") + broken["litellm_params"]["complexity_router_config"]["tiers"] = {} + assert router.upsert_deployment(Deployment(**broken)) is None + + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + live = router.get_deployment(model_id="id-a") + assert live is not None and live.litellm_params.complexity_router_config["classifier_type"] == "heuristic_v2" + assert router.heuristic_v2_router_limit_violation() is not None + + def test_heuristic_v2_routers_are_unlimited_by_default(self) -> None: + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ] + ) + + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + assert router.heuristic_v2_router_limit_violation() is None + + def test_heuristic_v2_router_limit_violation_frees_the_slot_of_the_router_being_edited(self) -> None: + """A DB reload upserts the existing heuristic_v2 router again; that edit must keep its own slot + while a different deployment switching to heuristic_v2 is refused.""" + router = Router( + model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")], + heuristic_v2_router_limit=lambda: 1, + ignore_invalid_deployments=True, + ) + + assert router.heuristic_v2_router_limit_violation() is not None + + edited = self._router_row("v2-a-renamed", "id-a", "heuristic_v2") + assert router.upsert_deployment(Deployment(**edited)) is not None + assert sorted(router.complexity_routers) == ["v2-a-renamed"] + + assert router.upsert_deployment(Deployment(**self._router_row("v2-b", "id-b", "heuristic_v2"))) is None + assert sorted(router.complexity_routers) == ["v2-a-renamed"] + assert router.upsert_deployment(Deployment(**self._router_row("v1-c", "id-c", "heuristic"))) is not None + assert sorted(router.complexity_routers) == ["v1-c", "v2-a-renamed"] + def test_hybrid_initialization_waits_for_later_pool_deployments(self): router = Router( model_list=[ diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py index 60433921de6..b5651062098 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -1,11 +1,12 @@ import asyncio +import itertools +import json +from collections.abc import Sequence +from typing import Final from unittest.mock import AsyncMock, patch import pytest - -import json - import litellm from litellm.caching.dual_cache import DualCache from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( @@ -102,11 +103,10 @@ async def test_async_user_key_affinity_routes_to_same_deployment(): # Deterministic routing: first selection uses seq[0], second selection attempts seq[1] # unless the list has been filtered to length=1 by deployment affinity. - choice_calls = {"count": 0} + choice_calls: Final = itertools.count(1) - def deterministic_choice(seq): - choice_calls["count"] += 1 - if choice_calls["count"] == 1: + def deterministic_choice(seq: Sequence[dict[str, object]]) -> dict[str, object]: + if next(choice_calls) == 1: return seq[0] return seq[1] if len(seq) > 1 else seq[0] @@ -998,3 +998,212 @@ async def test_model_group_affinity_config_overrides_global(): ) # All deployments returned (user-key affinity disabled for this group) assert len(filtered) == 2 + + +def _jwt_metadata(user_id: str) -> dict[str, str | None]: + return {"user_api_key_hash": None, "user_api_key_user_id": user_id} + + +def _two_deployments(model_group: str) -> list[dict]: + return [ + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "openai-deployment-a"}, + }, + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "openai-deployment-b"}, + }, + ] + + +@pytest.mark.asyncio +async def test_async_jwt_user_affinity_routes_to_same_deployment(): + """ + JWT-authenticated proxy requests carry no `user_api_key_hash`, only `user_api_key_user_id`. + They must still pin to one deployment per user. + """ + model_group = "gpt-5.4-mini" + router = litellm.Router( + model_list=[ + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "mock-api-key-a"}, + "model_info": {"id": "openai-deployment-a"}, + }, + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "mock-api-key-b"}, + "model_info": {"id": "openai-deployment-b"}, + }, + ], + optional_pre_call_checks=["deployment_affinity"], + ) + + choice_calls = {"count": 0} + + def deterministic_choice(seq): + choice_calls["count"] += 1 + if choice_calls["count"] == 1: + return seq[0] + return seq[1] if len(seq) > 1 else seq[0] + + with patch( # test-quality-ok: simple-shuffle has no injectable RNG; forcing the other pick is what proves the pin overrides the strategy + "litellm.router_strategy.simple_shuffle.random.choice", + side_effect=deterministic_choice, + ): + first_response = await router.acompletion( + model=model_group, + messages=[{"role": "user", "content": "Reply with the single word ok"}], + mock_response="ok", + metadata=_jwt_metadata("jwt-user-alice"), + ) + second_response = await router.acompletion( + model=model_group, + messages=[{"role": "user", "content": "Reply with the single word ok"}], + mock_response="ok", + metadata=_jwt_metadata("jwt-user-alice"), + ) + + first_model_id = first_response._hidden_params["model_id"] + assert first_model_id in ("openai-deployment-a", "openai-deployment-b") + assert second_response._hidden_params["model_id"] == first_model_id + + +@pytest.mark.asyncio +async def test_proxy_jwt_auth_metadata_pins_per_user(): + """ + The metadata the proxy stamps for a JWT caller (`UserAPIKeyAuth(api_key=None, user_id=)`) + must claim a pin and be read back by the filter, and another JWT user must not inherit it. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + model_group = "gpt-5.4-mini" + healthy_deployments = _two_deployments(model_group) + callback = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=60, + enable_user_key_affinity=True, + enable_responses_api_affinity=False, + ) + + def proxy_request(user_id: str) -> dict[str, object]: + return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data={"model": model_group, "messages": [{"role": "user", "content": "hi"}], "metadata": {}}, + user_api_key_dict=UserAPIKeyAuth(api_key=None, user_id=user_id), + _metadata_variable_name="metadata", + ) + + alice_request = proxy_request("jwt-user-alice") + alice_metadata = alice_request["metadata"] + assert isinstance(alice_metadata, dict) + assert alice_metadata["user_api_key_hash"] is None + + await callback.async_pre_call_deployment_hook( + kwargs={ + **alice_request, + "metadata": {**alice_metadata, "deployment_model_name": model_group}, + "model_info": {"id": "openai-deployment-b"}, + }, + call_type=None, + ) + + alice_pinned = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=alice_request, + parent_otel_span=None, + ) + assert [deployment["model_info"]["id"] for deployment in alice_pinned] == ["openai-deployment-b"] + + bob_filtered = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=proxy_request("jwt-user-bob"), + parent_otel_span=None, + ) + assert bob_filtered == healthy_deployments + + +@pytest.mark.asyncio +async def test_jwt_user_id_never_reads_a_virtual_key_pin(): + """ + A JWT user id that happens to equal a virtual key's 64-hex hash must not read that key's pin. + """ + model_group = "gpt-5.4-mini" + healthy_deployments = _two_deployments(model_group) + callback = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=60, + enable_user_key_affinity=True, + enable_responses_api_affinity=False, + ) + key_hash = "a" * 64 + + await callback.async_pre_call_deployment_hook( + kwargs={ + "metadata": {"user_api_key_hash": key_hash, "deployment_model_name": model_group}, + "model_info": {"id": "openai-deployment-b"}, + }, + call_type=None, + ) + + key_pinned = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={"metadata": {"user_api_key_hash": key_hash}}, + parent_otel_span=None, + ) + assert [deployment["model_info"]["id"] for deployment in key_pinned] == ["openai-deployment-b"] + + lookalike_jwt_user = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={"metadata": _jwt_metadata(key_hash)}, + parent_otel_span=None, + ) + assert lookalike_jwt_user == healthy_deployments + + +@pytest.mark.asyncio +async def test_virtual_key_hash_wins_over_user_id_for_affinity(): + """ + A virtual-key caller with a user id pins on the key hash, so two keys owned by one user + keep independent pins. + """ + model_group = "gpt-5.4-mini" + healthy_deployments = _two_deployments(model_group) + callback = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=60, + enable_user_key_affinity=True, + enable_responses_api_affinity=False, + ) + + await callback.async_pre_call_deployment_hook( + kwargs={ + "metadata": { + "user_api_key_hash": "key-one", + "user_api_key_user_id": "shared-user", + "deployment_model_name": model_group, + }, + "model_info": {"id": "openai-deployment-b"}, + }, + call_type=None, + ) + + other_key_same_user = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={"metadata": {"user_api_key_hash": "key-two", "user_api_key_user_id": "shared-user"}}, + parent_otel_span=None, + ) + assert other_key_same_user == healthy_deployments diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/test_litellm/router_utils/test_auto_router_model_naming.py index 0007f09896a..238d0546518 100644 --- a/tests/test_litellm/router_utils/test_auto_router_model_naming.py +++ b/tests/test_litellm/router_utils/test_auto_router_model_naming.py @@ -1,8 +1,13 @@ +from collections.abc import Mapping + import pytest from litellm.router_utils.auto_router_model_naming import ( carries_complexity_router_settings, classify_strategy_router_model, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, + is_heuristic_v2_router, strategy_router_dependencies, validate_complexity_router_config_placement, validate_complexity_router_config_write, @@ -369,3 +374,55 @@ def test_placement_is_scoped_to_complexity_router_deployments(model, present_fie flat param on an s3_vectors vector store, so an unscoped gate would reject a valid deployment. Either complexity field names one on its own, which is what the load itself requires.""" assert carries_complexity_router_settings(model, present_fields) is scoped + + +@pytest.mark.parametrize( + "litellm_params,expected", + [ + ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True), + ({"model": "auto_router/complexity_router-eu", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic"}}, False), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"tiers": {"SIMPLE": "a"}}}, False), + ({"model": "auto_router/complexity_router"}, False), + ({"model": "auto_router/quality_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False), + ({"model": "openai/gpt-4o", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False), + ({"model": "auto_router/complexity_router", "complexity_router_config": "heuristic_v2"}, False), + ({}, False), + ], +) +def test_is_heuristic_v2_router(litellm_params: Mapping[str, object], expected: bool) -> None: + """Only a complexity router whose config selects heuristic_v2 counts toward the license limit.""" + assert is_heuristic_v2_router(litellm_params) is expected + + +def test_count_heuristic_v2_routers_reads_model_list_rows_and_ignores_malformed_ones() -> None: + v2 = {"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}} + rows: list[Mapping[str, object]] = [ + {"model_name": "a", "litellm_params": v2}, + {"model_name": "b", "litellm_params": {"model": "openai/gpt-4o"}}, + {"model_name": "c", "litellm_params": v2}, + {"model_name": "d"}, + {"model_name": "e", "litellm_params": "not a mapping"}, + ] + assert count_heuristic_v2_routers(rows) == 2 + assert count_heuristic_v2_routers(()) == 0 + + +@pytest.mark.parametrize( + "held,limit,violates", + [ + (1, 1, False), + (2, 1, True), + (0, 1, False), + (5, None, False), + (3, 3, False), + (4, 3, True), + ], +) +def test_heuristic_v2_limit_violation(held: int, limit: int | None, violates: bool) -> None: + violation = heuristic_v2_limit_violation(held=held, limit=limit) + assert (violation is not None) is violates + if violation is not None: + assert f"At most {limit} auto-router" in violation + assert f"would make {held}" in violation + assert "license" not in violation diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index a0dbf3b6637..7b2e45ab3ed 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -371,3 +371,20 @@ class TestKimiK3AdvertisesItsDocumentedLevels: "low", "high", ) + + +class TestGpt6AstraAdvertisesItsDocumentedLevels: + def test_the_entry_advertises_low_through_max_without_none(self, local_model_cost_map): + """OpenAI documents low, medium, high, xhigh and max for gpt-6-astra. Unlike gpt-5.6-sol it + does not take none, so a group must not offer none and must offer max.""" + from litellm.utils import _get_model_info_helper + + model_info = dict(_get_model_info_helper(model="gpt-6-astra", custom_llm_provider="openai")) + + assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( + "low", + "medium", + "high", + "xhigh", + "max", + ) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 862d77f01b2..2046695f151 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4478,3 +4478,17 @@ def test_explicit_pricing_precedes_private_provider_response_model( ) assert selected == expected + + +def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_model_cost_map): + """gpt-6-astra batch pricing is 50% off the standard $10 input and $50 output rates per 1M tokens.""" + from litellm.cost_calculator import batch_cost_calculator + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + prompt_cost, completion_cost = batch_cost_calculator( + usage=usage, model="gpt-6-astra", custom_llm_provider="openai" + ) + + assert prompt_cost == pytest.approx(1000 * 5e-6) + assert completion_cost == pytest.approx(500 * 2.5e-5) diff --git a/tests/test_litellm/test_gpt_5_5_model_metadata.py b/tests/test_litellm/test_gpt_5_5_model_metadata.py index 1c12a48ed9d..a60fa9466e6 100644 --- a/tests/test_litellm/test_gpt_5_5_model_metadata.py +++ b/tests/test_litellm/test_gpt_5_5_model_metadata.py @@ -47,8 +47,7 @@ def test_azure_ai_gpt_5_5_model_info(model): routed_model, provider, _, _ = get_llm_provider(model=model) assert routed_model == model.split("/", 1)[1] - # azure_ai/* models resolve under the azure provider in get_llm_provider - assert provider == "azure" + assert provider == "azure_ai" def test_azure_ai_gpt_5_5_backup_matches_main(): diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 3c8bf142835..9a703635ab6 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -784,6 +784,19 @@ def test_responses_api_bridge_check_gpt_5_4_tools_plus_reasoning_routes_to_respo assert model_info.get("mode") == "responses" +def test_responses_api_bridge_check_gpt_6_astra_tools_with_default_reasoning_routes_to_responses(): + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model="gpt-6-astra", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + ) + + assert model == "gpt-6-astra" + assert model_info.get("mode") == "responses" + + def test_responses_api_bridge_check_gpt_5_5_tools_plus_reasoning_routes_to_responses(): """gpt-5.5+ with both tools and reasoning_effort should route to Responses API.""" from litellm.main import responses_api_bridge_check @@ -1133,6 +1146,82 @@ def test_responses_api_bridge_check_custom_api_base_via_env_with_unset_effort_st assert model_info.get("mode") != "responses" +@pytest.mark.parametrize( + "api_base", + [ + "https://southcentralus.privatelink.api.openai.com/v1", + "https://privatelink.corp.api.openai.com/v1", + "https://api.openai.com:443/v1", + "https://api.openai.com/v1/", + "HTTPS://API.OPENAI.COM/v1", + ], +) +def test_responses_api_bridge_check_openai_backed_custom_api_base_with_unset_effort_routes_to_responses(api_base): + """ + A custom api_base whose host is api.openai.com or a subdomain of it (a PrivateLink hostname, a + port-qualified or trailing-slash default) still reaches the real OpenAI backend, which rejects + function tools with reasoning on Chat Completions, so the unset-effort arm must bridge exactly as + it does for the literal default URL. Regression guard for GH #39353. + """ + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=api_base, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +@pytest.mark.parametrize( + "api_base", + [ + "https://api.openai.com.evil.example/v1", + "https://notapi.openai.com/v1", + "https://gateway.example/v1?upstream=api.openai.com", + "https://openai.internal.example/api.openai.com/v1", + ], +) +def test_responses_api_bridge_check_lookalike_custom_api_base_with_unset_effort_stays_chat(api_base): + """Only the host decides: api.openai.com appearing elsewhere in the URL is still a foreign backend.""" + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=api_base, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_privatelink_api_base_via_env_with_unset_effort_routes_to_responses(monkeypatch): + """A PrivateLink base set through OPENAI_BASE_URL resolves the way the chat handler's does and still bridges.""" + import litellm + from litellm.main import responses_api_bridge_check + + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setenv("OPENAI_BASE_URL", "https://southcentralus.privatelink.api.openai.com/v1") + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=None, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + def test_responses_api_bridge_check_custom_api_base_with_explicit_effort_still_routes(): """Explicit reasoning_effort keeps its pre-existing bridging behavior on any api_base.""" from litellm.main import responses_api_bridge_check @@ -3259,3 +3348,43 @@ def test_stream_chunk_builder_leaves_xai_reported_cost_to_the_calculator(monkeyp assert getattr(response.usage, "cost", None) == pytest.approx(0.42) assert response._hidden_params.get("response_cost") is None assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.63) + + +FOUNDRY_HOST: Final = "https://my-project.services.ai.azure.com" + + +def test_azure_ai_transcription_on_a_foundry_host_uses_the_azure_openai_deployment_route( + respx_mock: respx.MockRouter, +): + route: Final = respx_mock.post( + url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/whisper-1/audio/transcriptions\?api-version=.+" + ).mock(return_value=httpx.Response(200, json={"text": "hello"})) + + response: Final = litellm.transcription( + model="azure_ai/whisper-1", + file=("tone.wav", b"RIFF\x00\x00\x00\x00WAVE", "audio/wav"), + api_base=FOUNDRY_HOST, + api_key="fake-key", + ) + + assert route.called + assert response.text == "hello" + + +def test_azure_ai_speech_on_a_foundry_host_uses_the_azure_openai_deployment_route( + respx_mock: respx.MockRouter, +): + route: Final = respx_mock.post( + url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/tts-1/audio/speech\?api-version=.+" + ).mock(return_value=httpx.Response(200, content=b"mp3-bytes")) + + response: Final = litellm.speech( + model="azure_ai/tts-1", + input="hello", + voice="alloy", + api_base=FOUNDRY_HOST, + api_key="fake-key", + ) + + assert route.called + assert response.content == b"mp3-bytes" diff --git a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py index c0860a5b55f..bdb2dc26813 100644 --- a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py +++ b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py @@ -52,6 +52,12 @@ PRIORITY_LONG_CONTEXT = { "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, }, + "gpt-6-astra": { + "input_cost_per_token_above_272k_tokens_priority": 4e-05, + "output_cost_per_token_above_272k_tokens_priority": 0.00015, + "cache_read_input_token_cost_above_272k_tokens_priority": 4e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, + }, } EXPECTED = {**FLEX_LONG_CONTEXT, **PRIORITY_LONG_CONTEXT} @@ -63,6 +69,7 @@ NO_PUBLISHED_PRIORITY_LONG_CONTEXT = ("gpt-5.4", "gpt-5.5") def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.add_known_models() @lru_cache(maxsize=2) @@ -114,6 +121,7 @@ TIERED_COST_CASES = [ ("gpt-5.6-sol", "priority", 1.6e-05, 6e-05), ("gpt-5.6-terra", "priority", 8e-06, 3.6e-05), ("gpt-5.6-luna", "priority", 8e-07, 3.6e-06), + ("gpt-6-astra", "priority", 4e-05, 0.00015), ] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 3b4c80b6b4f..26c0209bc05 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -12553,3 +12553,33 @@ async def test_prompt_management_factory_marks_injection_for_every_deployment(mo bucket = captured.get("litellm_metadata") or captured["metadata"] assert captured["model_info"]["id"] == "provisional-dep" assert bucket["litellm_gateway_injected_cache"] == "" + + +def test_get_configured_mode_reads_deployment_model_info(): + router = Router( + model_list=[ + { + "model_name": "my-tts", + "litellm_params": {"model": "openai/some-unmapped-mode-model"}, + "model_info": {"mode": "audio_speech"}, + } + ] + ) + + assert router.get_configured_mode("my-tts") == "audio_speech" + + +@pytest.mark.parametrize("model_info", [{}, {"mode": ""}, {"mode": " "}, {"mode": 123}]) +def test_get_configured_mode_returns_none_for_unset_blank_or_unknown(model_info): + router = Router( + model_list=[ + { + "model_name": "plain-model", + "litellm_params": {"model": "openai/some-unmapped-mode-model"}, + "model_info": model_info, + } + ] + ) + + assert router.get_configured_mode("plain-model") is None + assert router.get_configured_mode("unknown-model") is None diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx index e0b2b897b36..524ec6a715e 100644 --- a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx @@ -58,6 +58,17 @@ describe("OrganizationDropdown", () => { expect(onChange.mock.calls[0][0]).toBe("org-1"); }); + it("emits null, never the empty string, when the selection is cleared", async () => { + const onChange = vi.fn(); + const user = userEvent.setup(); + render(); + + await user.click(screen.getByRole("button", { name: "Clear" })); + + expect(onChange).toHaveBeenCalledTimes(1); + expect(onChange).toHaveBeenCalledWith(null); + }); + it("should filter options by organization id", async () => { const user = userEvent.setup(); render(); diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx index 8da35ecd02e..663028b2d92 100644 --- a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx @@ -5,7 +5,7 @@ import { Organization } from "../networking"; interface OrganizationDropdownProps { organizations?: Organization[] | null; value?: string; - onChange?: (value: string) => void; + onChange?: (value: string | null) => void; disabled?: boolean; loading?: boolean; style?: React.CSSProperties; @@ -32,7 +32,7 @@ const OrganizationDropdown: React.FC = ({ sublabel: org.organization_id, }))} value={value} - onValueChange={(organizationId) => onChange?.(organizationId)} + onValueChange={(organizationId) => onChange?.(organizationId || null)} placeholder={placeholder} emptyText={loading ? "Loading organizations…" : "No organizations found"} disabled={disabled} diff --git a/ui/litellm-dashboard/src/components/object_permissions_view.tsx b/ui/litellm-dashboard/src/components/object_permissions_view.tsx index 011791fa18d..e90645a9e9f 100644 --- a/ui/litellm-dashboard/src/components/object_permissions_view.tsx +++ b/ui/litellm-dashboard/src/components/object_permissions_view.tsx @@ -3,9 +3,12 @@ import VectorStorePermissions from "./permissions/VectorStorePermissions"; import MCPServerPermissions from "./permissions/MCPServerPermissions"; import AgentPermissions from "./permissions/AgentPermissions"; import type { ObjectPermission } from "./object_permission_types"; +import type { InheritedGrant } from "./permissions/inheritedGrants"; interface ObjectPermissionsViewProps { objectPermission?: ObjectPermission | null; + inheritedMcpServers?: InheritedGrant[]; + inheritedAgents?: InheritedGrant[]; variant?: "card" | "inline"; className?: string; accessToken?: string | null; @@ -13,6 +16,8 @@ interface ObjectPermissionsViewProps { export function ObjectPermissionsView({ objectPermission, + inheritedMcpServers = [], + inheritedAgents = [], variant = "card", className = "", accessToken, @@ -34,9 +39,15 @@ export function ObjectPermissionsView({ mcpAccessGroups={mcpAccessGroups} mcpToolPermissions={mcpToolPermissions} mcpToolsets={mcpToolsets} + inheritedMcpServers={inheritedMcpServers} + accessToken={accessToken} + /> + -

Search tools

{searchTools.length === 0 ? ( diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 5749541dcea..ee3a88acba0 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -587,9 +587,9 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp } }; - const changeOrganization = (write: FieldWrite) => (orgId: string) => { - write(orgId || undefined); - setSelectedOrganizationId(orgId || null); + const changeOrganization = (write: FieldWrite) => (orgId: string | null) => { + write(orgId ?? undefined); + setSelectedOrganizationId(orgId); // Clear team and project when org changes setSelectedCreateKeyTeam(null); setSelectedProjectId(null); diff --git a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx new file mode 100644 index 00000000000..01c7b60b8a5 --- /dev/null +++ b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx @@ -0,0 +1,71 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import AgentPermissions from "./AgentPermissions"; +import * as networking from "../networking"; + +vi.mock("../networking"); + +describe("AgentPermissions", () => { + const accessToken = "test-token"; + const agentId = "90337622-756e-4f25-98f0-01fc8174aa24"; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("lists agents inherited from access groups, counts them, and names the groups on hover", async () => { + const user = userEvent.setup(); + vi.mocked(networking.getAgentsList).mockResolvedValue({ + agents: [{ agent_id: agentId, agent_name: "support_agent" }], + }); + + render( + , + ); + + const row = await screen.findByText(/support_agent/); + expect(screen.getByText("1")).toBeInTheDocument(); + expect(screen.queryByText("No agents or access groups configured")).not.toBeInTheDocument(); + expect(networking.getAgentsList).toHaveBeenCalledWith(accessToken); + + await user.hover(row); + expect( + await screen.findByText(`Granted via access groups platform-tools, support. Full ID: ${agentId}`), + ).toBeInTheDocument(); + }); + + it("does not double-list an agent that is both granted directly and inherited", async () => { + const user = userEvent.setup(); + vi.mocked(networking.getAgentsList).mockResolvedValue({ + agents: [{ agent_id: agentId, agent_name: "support_agent" }], + }); + + render( + , + ); + + const row = await screen.findByText(/support_agent/); + expect(screen.getAllByText(/support_agent/)).toHaveLength(1); + expect(screen.getByText("1")).toBeInTheDocument(); + + await user.hover(row); + expect(await screen.findByText(`Full ID: ${agentId}`)).toBeInTheDocument(); + }); + + it("shows the empty state when nothing is granted directly or inherited", () => { + render(); + + expect(screen.getByText("No agents or access groups configured")).toBeInTheDocument(); + expect(screen.getByText("0")).toBeInTheDocument(); + expect(networking.getAgentsList).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx index d84bd1834f8..d1ca25c975b 100644 --- a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx @@ -3,6 +3,7 @@ import { UserGroupIcon } from "@heroicons/react/outline"; import { Badge } from "@/components/ui/badge"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { getAgentsList } from "../networking"; +import { InheritedGrant, inheritedGrantTooltip } from "./inheritedGrants"; interface Agent { agent_id: string; @@ -14,16 +15,24 @@ interface Agent { interface AgentPermissionsProps { agents: string[]; agentAccessGroups?: string[]; + inheritedAgents?: InheritedGrant[]; accessToken?: string | null; } -export function AgentPermissions({ agents, agentAccessGroups = [], accessToken }: AgentPermissionsProps) { +export function AgentPermissions({ + agents, + agentAccessGroups = [], + inheritedAgents = [], + accessToken, +}: AgentPermissionsProps) { const [agentDetails, setAgentDetails] = useState([]); + const inheritedOnlyAgents = inheritedAgents.filter((grant) => !agents.includes(grant.id)); + const agentIdCount = agents.length + inheritedOnlyAgents.length; // Fetch agent details when component mounts useEffect(() => { const fetchAgentDetails = async () => { - if (accessToken && agents.length > 0) { + if (accessToken && agentIdCount > 0) { try { const response = await getAgentsList(accessToken); if (response && response.agents && Array.isArray(response.agents)) { @@ -35,7 +44,7 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken } } }; fetchAgentDetails(); - }, [accessToken, agents.length]); + }, [accessToken, agentIdCount]); // Function to get display name for agent const getAgentDisplayName = (agentId: string) => { @@ -47,10 +56,10 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken } return agentId; }; - // Merge agents and access groups into one list const mergedItems = [ - ...agents.map((agent) => ({ type: "agent", value: agent })), - ...agentAccessGroups.map((group) => ({ type: "accessGroup", value: group })), + ...agents.map((agent) => ({ type: "agent", value: agent, tooltip: `Full ID: ${agent}` })), + ...inheritedOnlyAgents.map((grant) => ({ type: "agent", value: grant.id, tooltip: inheritedGrantTooltip(grant) })), + ...agentAccessGroups.map((group) => ({ type: "accessGroup", value: group, tooltip: "" })), ]; const totalCount = mergedItems.length; @@ -77,7 +86,7 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken } {getAgentDisplayName(item.value)} - {`Full ID: ${item.value}`} + {item.tooltip} ) : ( diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx index c2df945f367..78ede9865a4 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx @@ -406,4 +406,55 @@ describe("MCPServerPermissions", () => { ); await waitFor(() => expect(screen.getByText("Blocked")).toHaveAttribute("data-variant", "destructive")); }); + + it("lists servers inherited from access groups, counts them, and names the group on hover", async () => { + const user = userEvent.setup(); + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { server_id: mockServerId1, server_name: mockServerName1, alias: mockServerName1 }, + ]); + + render( + , + ); + + const row = await screen.findByText(/DW_MCP/); + expect(screen.getByText("1")).toBeInTheDocument(); + expect(screen.queryByText("No MCP servers, access groups, or toolsets configured")).not.toBeInTheDocument(); + expect(networking.fetchMCPServers).toHaveBeenCalledWith(mockAccessToken); + + await user.hover(row); + expect( + await screen.findByText(`Granted via access group platform-tools. Full ID: ${mockServerId1}`), + ).toBeInTheDocument(); + }); + + it("does not double-list a server that is both granted directly and inherited", async () => { + const user = userEvent.setup(); + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { server_id: mockServerId2, server_name: mockServerName2, alias: mockServerName2 }, + ]); + + render( + , + ); + + const row = await screen.findByText(/Test Server/); + expect(screen.getAllByText(/Test Server/)).toHaveLength(1); + expect(screen.getByText("1")).toBeInTheDocument(); + + await user.hover(row); + expect(await screen.findByText(`Full ID: ${mockServerId2}`)).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx index f11bc6b9627..65f210addfa 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx @@ -5,12 +5,14 @@ import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip import { fetchMCPServers, fetchMCPToolsets } from "../networking"; import { MCPServer, MCPToolset } from "../mcp_tools/types"; import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; +import { InheritedGrant, inheritedGrantTooltip } from "./inheritedGrants"; interface MCPServerPermissionsProps { mcpServers: string[]; mcpAccessGroups?: string[]; mcpToolPermissions?: Record; mcpToolsets?: string[]; + inheritedMcpServers?: InheritedGrant[]; accessToken?: string | null; } @@ -19,6 +21,7 @@ export function MCPServerPermissions({ mcpAccessGroups = [], mcpToolPermissions = {}, mcpToolsets = [], + inheritedMcpServers = [], accessToken, }: MCPServerPermissionsProps) { const [mcpServerDetails, setMCPServerDetails] = useState([]); @@ -50,10 +53,16 @@ export function MCPServerPermissions({ }); }; + const directServerIds = mcpServers.filter( + (server) => server !== NO_MCP_SERVERS_SENTINEL && server !== ALL_PROXY_MCP_SERVERS_SENTINEL, + ); + const inheritedOnlyServers = inheritedMcpServers.filter((grant) => !mcpServers.includes(grant.id)); + const serverIdCount = directServerIds.length + inheritedOnlyServers.length; + // Fetch MCP server details when component mounts useEffect(() => { const fetchMCPServerDetails = async () => { - if (accessToken && mcpServers.length > 0) { + if (accessToken && serverIdCount > 0) { try { const response = await fetchMCPServers(accessToken); if (response && Array.isArray(response)) { @@ -67,7 +76,7 @@ export function MCPServerPermissions({ } }; fetchMCPServerDetails(); - }, [accessToken, mcpServers.length]); + }, [accessToken, serverIdCount]); // Fetch toolset details useEffect(() => { @@ -98,12 +107,14 @@ export function MCPServerPermissions({ const blocksAllMcpServers = mcpServers.includes(NO_MCP_SERVERS_SENTINEL); const grantsAllProxyMcpServers = mcpServers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL); - // Merge servers and access groups into one list const mergedItems = [ - ...mcpServers - .filter((server) => server !== NO_MCP_SERVERS_SENTINEL && server !== ALL_PROXY_MCP_SERVERS_SENTINEL) - .map((server) => ({ type: "server", value: server })), - ...mcpAccessGroups.map((group) => ({ type: "accessGroup", value: group })), + ...directServerIds.map((server) => ({ type: "server", value: server, tooltip: `Full ID: ${server}` })), + ...inheritedOnlyServers.map((grant) => ({ + type: "server", + value: grant.id, + tooltip: inheritedGrantTooltip(grant), + })), + ...mcpAccessGroups.map((group) => ({ type: "accessGroup", value: group, tooltip: "" })), ]; const totalCount = mergedItems.length + mcpToolsets.length; @@ -153,7 +164,7 @@ export function MCPServerPermissions({ {getMCPServerDisplayName(item.value)} - {`Full ID: ${item.value}`} + {item.tooltip} ) : (
diff --git a/ui/litellm-dashboard/src/components/permissions/inheritedGrants.test.ts b/ui/litellm-dashboard/src/components/permissions/inheritedGrants.test.ts new file mode 100644 index 00000000000..766f8f23f0b --- /dev/null +++ b/ui/litellm-dashboard/src/components/permissions/inheritedGrants.test.ts @@ -0,0 +1,61 @@ +import { describe, expect, it } from "vitest"; +import { computeInheritedGrants, inheritedGrantTooltip } from "./inheritedGrants"; +import { TeamAccessGroupModelGrant } from "../team/teamModelAccess"; + +const GRANTS: TeamAccessGroupModelGrant[] = [ + { access_group_id: "ag-1", access_group_name: "platform-tools", models: [], mcp_server_ids: ["mcp-1", "mcp-2"] }, + { + access_group_id: "ag-2", + access_group_name: "support", + models: [], + mcp_server_ids: ["mcp-2"], + agent_ids: ["agent-1"], + }, +]; + +describe("computeInheritedGrants", () => { + it("attributes each id to every group that grants it, in group order", () => { + expect(computeInheritedGrants(["mcp-1", "mcp-2"], GRANTS, (g) => g.mcp_server_ids)).toEqual([ + { id: "mcp-1", accessGroupNames: ["platform-tools"] }, + { id: "mcp-2", accessGroupNames: ["platform-tools", "support"] }, + ]); + }); + + it("keeps ids the flat list carries but no group detail explains, with no group names", () => { + expect(computeInheritedGrants(["agent-1", "agent-legacy"], GRANTS, (g) => g.agent_ids)).toEqual([ + { id: "agent-1", accessGroupNames: ["support"] }, + { id: "agent-legacy", accessGroupNames: [] }, + ]); + }); + + it("falls back to the group details when the flat list is missing, without duplicates", () => { + expect(computeInheritedGrants(undefined, GRANTS, (g) => g.mcp_server_ids).map((g) => g.id)).toEqual([ + "mcp-1", + "mcp-2", + ]); + }); + + it("returns nothing when neither source has ids", () => { + expect(computeInheritedGrants(undefined, undefined, (g) => g.agent_ids)).toEqual([]); + }); +}); + +describe("inheritedGrantTooltip", () => { + it("names a single group", () => { + expect(inheritedGrantTooltip({ id: "mcp-1", accessGroupNames: ["platform-tools"] })).toBe( + "Granted via access group platform-tools. Full ID: mcp-1", + ); + }); + + it("lists several groups", () => { + expect(inheritedGrantTooltip({ id: "mcp-2", accessGroupNames: ["platform-tools", "support"] })).toBe( + "Granted via access groups platform-tools, support. Full ID: mcp-2", + ); + }); + + it("stays generic when the proxy did not say which group granted it", () => { + expect(inheritedGrantTooltip({ id: "agent-legacy", accessGroupNames: [] })).toBe( + "Granted via an access group. Full ID: agent-legacy", + ); + }); +}); diff --git a/ui/litellm-dashboard/src/components/permissions/inheritedGrants.ts b/ui/litellm-dashboard/src/components/permissions/inheritedGrants.ts new file mode 100644 index 00000000000..fb78de5ed77 --- /dev/null +++ b/ui/litellm-dashboard/src/components/permissions/inheritedGrants.ts @@ -0,0 +1,26 @@ +import { describeGroups, TeamAccessGroupModelGrant } from "../team/teamModelAccess"; + +export interface InheritedGrant { + id: string; + accessGroupNames: string[]; +} + +export function computeInheritedGrants( + ids: string[] | undefined, + grants: TeamAccessGroupModelGrant[] | undefined, + idsOf: (grant: TeamAccessGroupModelGrant) => string[] | undefined, +): InheritedGrant[] { + const known = grants ?? []; + const allIds = [...new Set([...(ids ?? []), ...known.flatMap((grant) => idsOf(grant) ?? [])])]; + return allIds.map((id) => ({ + id, + accessGroupNames: known + .filter((grant) => (idsOf(grant) ?? []).includes(id)) + .map((grant) => grant.access_group_name), + })); +} + +export const inheritedGrantTooltip = (grant: InheritedGrant): string => { + const source = grant.accessGroupNames.length > 0 ? describeGroups(grant.accessGroupNames) : "an access group"; + return `Granted via ${source}. Full ID: ${grant.id}`; +}; diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index ca1e0413dcb..aedc04283f5 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -38,6 +38,9 @@ vi.mock("@/components/networking", () => ({ organizationInfoCall: vi.fn(), getRouterSettingsCall: vi.fn().mockResolvedValue({ fields: [] }), getPassThroughEndpointsCall: vi.fn(), + fetchMCPServers: vi.fn().mockResolvedValue([]), + fetchMCPToolsets: vi.fn().mockResolvedValue([]), + getAgentsList: vi.fn().mockResolvedValue({ agents: [] }), })); const can = vi.fn(); @@ -302,6 +305,47 @@ describe("TeamInfoView", () => { ); }); + it("shows MCP servers and agents inherited from access groups in the Object Permissions card, naming the group on hover", async () => { + const user = userEvent.setup(); + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { server_id: "mcp-github-1234", server_name: "github", alias: "github" }, + ]); + vi.mocked(networking.getAgentsList).mockResolvedValue({ + agents: [{ agent_id: "agent-support-5678", agent_name: "support_agent" }], + }); + const platformToolsGroup = { + access_group_id: "ag-1", + access_group_name: "platform-tools", + models: [], + mcp_server_ids: ["mcp-github-1234"], + agent_ids: ["agent-support-5678"], + }; + const inheritedGrants = { + object_permission: null, + access_group_ids: ["ag-1"], + access_group_mcp_server_ids: ["mcp-github-1234"], + access_group_agent_ids: ["agent-support-5678"], + access_group_details: [platformToolsGroup], + }; + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData(inheritedGrants)); + + renderWithProviders(); + + const serverRow = await screen.findByText(/github \(mcp\.\.\.1234\)/); + const agentRow = await screen.findByText(/support_agent \(age\.\.\.5678\)/); + expect(screen.queryByText("No MCP servers, access groups, or toolsets configured")).not.toBeInTheDocument(); + expect(screen.queryByText("No agents or access groups configured")).not.toBeInTheDocument(); + + await user.hover(serverRow); + expect( + await screen.findByText("Granted via access group platform-tools. Full ID: mcp-github-1234"), + ).toBeInTheDocument(); + await user.hover(agentRow); + expect( + await screen.findByText("Granted via access group platform-tools. Full ID: agent-support-5678"), + ).toBeInTheDocument(); + }); + it("keeps the all-proxy-models badge non-clickable", async () => { vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["all-proxy-models"] })); @@ -1844,6 +1888,8 @@ describe("TeamInfoView - the exact bytes the update call sends", () => { mcp_access_groups: [], mcp_tool_permissions: {}, mcp_toolsets: [], + agents: [], + agent_access_groups: [], vector_stores: ["vs-1"], }; @@ -1864,6 +1910,49 @@ describe("TeamInfoView - the exact bytes the update call sends", () => { }); }); + const openEditorWithAgents = async (user: ReturnType) => { + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + models: ["gpt-4"], + object_permission: { agents: ["agent-1"], agent_access_groups: ["group-a"] }, + }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0)); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await screen.findByLabelText("Team Name"); + }; + + it("resends the stored agents and agent_access_groups when the selector is left untouched", async () => { + const user = userEvent.setup({ delay: null }); + await openEditorWithAgents(user); + + const payload = await save(user); + + const objectPermission = wireBody(payload).object_permission as Record; + expect(objectPermission.agents).toStrictEqual(["agent-1"]); + expect(objectPermission.agent_access_groups).toStrictEqual(["group-a"]); + }); + + it("sends empty agents and agent_access_groups arrays after the last agent chip is removed", async () => { + const user = userEvent.setup({ delay: null }); + await openEditorWithAgents(user); + + await user.click(within(screen.getByLabelText("agent-1")).getByRole("button")); + await user.click(within(screen.getByLabelText("group:group-a")).getByRole("button")); + expect(screen.queryByLabelText("agent-1")).not.toBeInTheDocument(); + expect(screen.queryByLabelText("group:group-a")).not.toBeInTheDocument(); + + const payload = await save(user); + + const objectPermission = wireBody(payload).object_permission as Record; + expect(objectPermission.agents).toStrictEqual([]); + expect(objectPermission.agent_access_groups).toStrictEqual([]); + }); + it("resends every stored value once both sections are opened", async () => { const user = userEvent.setup({ delay: null }); await openEditor(user); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 44f0d420fd6..c2b8cd3cc56 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -58,6 +58,7 @@ import { TeamModelBadge, TeamModelBadgeKind, } from "./teamModelAccess"; +import { computeInheritedGrants } from "../permissions/inheritedGrants"; import MetadataKeyValueFields, { metadataObjectToPairs, metadataPairsSchema, @@ -862,12 +863,8 @@ const TeamInfoView: React.FC = ({ agents: [], accessGroups: [], }; - if (agents && agents.length > 0) { - updateData.object_permission.agents = agents; - } - if (agentAccessGroups && agentAccessGroups.length > 0) { - updateData.object_permission.agent_access_groups = agentAccessGroups; - } + updateData.object_permission.agents = agents; + updateData.object_permission.agent_access_groups = agentAccessGroups; delete values.agents_and_groups; // Handle vector stores permissions @@ -936,6 +933,17 @@ const TeamInfoView: React.FC = ({ const { team_info: info } = teamData; + const inheritedMcpServers = computeInheritedGrants( + info.access_group_mcp_server_ids, + info.access_group_details, + (grant) => grant.mcp_server_ids, + ); + const inheritedAgents = computeInheritedGrants( + info.access_group_agent_ids, + info.access_group_details, + (grant) => grant.agent_ids, + ); + const initialKillSwitchOn = info.metadata?.disable_global_guardrails === true; const allGuardrails: GuardrailListItem[] = guardrailsData?.guardrails ?? []; @@ -1033,7 +1041,13 @@ const TeamInfoView: React.FC = ({
- + = ({ 0 ? models : [NO_DEFAULT_MODELS]; } -const describeGroups = (names: string[]): string => +export const describeGroups = (names: string[]): string => names.length > 1 ? `access groups ${names.join(", ")}` : `access group ${names[0]}`; export function computeTeamModelBadges( diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 3e772fd0e9b..1330be2788b 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -303,9 +303,9 @@ export function KeyEditView({ } }; - const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | undefined) => { - setField(orgId || null); - setSelectedOrganizationId(orgId || null); + const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | null) => { + setField(orgId); + setSelectedOrganizationId(orgId); form.setValue("team_id", undefined); }; diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx index e3bacc0908a..3c9e6543c1c 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx @@ -75,6 +75,41 @@ describe("Cost column", () => { }); }); +describe("Tokens column", () => { + const sessionRow: Partial = { + request_id: "req-session-tokens", + total_tokens: 10, + prompt_tokens: 7, + completion_tokens: 3, + session_id: "sess-1", + session_total_count: 3, + }; + + it("shows the summed session token usage, not the representative call's tokens, for a multi-round session", () => { + const aggregatedRow: Partial = { + ...sessionRow, + session_total_tokens: 60, + session_total_prompt_tokens: 42, + session_total_completion_tokens: 18, + }; + renderRows([logEntry(aggregatedRow)]); + + const tokensCell = screen.getByRole("cell", { name: /\(42\+18\)/ }); + expect(tokensCell).toHaveTextContent("60"); + expect(tokensCell).toHaveTextContent("session total"); + expect(screen.queryByText("10")).not.toBeInTheDocument(); + expect(screen.queryByText("(7+3)")).not.toBeInTheDocument(); + }); + + it("falls back to the call's own tokens with no session label when the backend sent no session token sums", () => { + renderRows([logEntry(sessionRow)]); + + const tokensCell = screen.getByRole("cell", { name: /\(7\+3\)/ }); + expect(tokensCell).toHaveTextContent("10"); + expect(tokensCell).not.toHaveTextContent("session total"); + }); +}); + describe("Type column", () => { it("shows the conversation badge and composition even when an MCP call represents the conversation", async () => { const user = userEvent.setup(); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx index 8db0b106851..9d4dc4f7898 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx @@ -263,13 +263,20 @@ export const getRequestLogsTableColumns = ({ meta: { numeric: true }, cell: ({ row }) => { const log = row.original; + const showSessionTotal = (log.session_total_count || 1) > 1 && log.session_total_tokens != null; + const total = showSessionTotal ? log.session_total_tokens : log.total_tokens; + const prompt = showSessionTotal ? log.session_total_prompt_tokens : log.prompt_tokens; + const completion = showSessionTotal ? log.session_total_completion_tokens : log.completion_tokens; return ( - - {String(log.total_tokens || "0")} - - ({String(log.prompt_tokens || "0")}+{String(log.completion_tokens || "0")}) +
+ + {String(total || "0")} + + ({String(prompt || "0")}+{String(completion || "0")}) + - + {showSessionTotal && session total} +
); }, }, diff --git a/ui/litellm-dashboard/src/components/view_logs/columns.tsx b/ui/litellm-dashboard/src/components/view_logs/columns.tsx index 21e09faf454..b2e29c3a0c1 100644 --- a/ui/litellm-dashboard/src/components/view_logs/columns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/columns.tsx @@ -42,6 +42,9 @@ export type LogEntry = { request_duration_ms?: number; session_total_count?: number; session_total_spend?: number; + session_total_tokens?: number; + session_total_prompt_tokens?: number; + session_total_completion_tokens?: number; session_cache_hit_count?: number; mcp_tool_call_count?: number; mcp_tool_call_spend?: number;