diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 4e428d8cebf..e85a397cbd2 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,7 +1,10 @@ + + ## TLDR - + Problem this solves: @@ -110,8 +113,20 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac ## Caveats (if any) - ## QA runbook @@ -134,6 +149,6 @@ Example checklists: - [ ] Sanity check: this test makes sense to add and is not hand-wavey (e.g., assert actual expected spend instead of just spend > 0) or potentially flaky --> -### Final Attestation +## Final Attestation - [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 2dfca3d308f..c23678c51ae 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -164,6 +164,7 @@ jobs: tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/rag_endpoints + tests/test_litellm/proxy/rerank_endpoints tests/test_litellm/proxy/realtime_endpoints tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/config_resolvers diff --git a/CLAUDE.md b/CLAUDE.md index b7bd06954d0..cbc6bd6848c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -37,13 +37,14 @@ If you're resolving a linear ticket, in the "## Linear ticket" section of the PR Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR -If you ever make public-facing PR descriptions, comments, issues, commit messages, etc., always follow these guidelines to sound less AI-y: +If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y: - don't use emojis - don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message - don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc. -- don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose +- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose - don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "." - don't use →. Instead, prefer not to use arrows, and if need be, use -> instead +- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 664e1669834..19c28716025 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 19955 + "limit": 19949 }, "reportArgumentType": { "limit": 2566 @@ -54,7 +54,7 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5663 + "limit": 5661 }, "reportMissingTypeArgument": { "limit": 15555 @@ -84,7 +84,7 @@ "limit": 56 }, "reportPrivateUsage": { - "limit": 1822 + "limit": 1810 }, "reportRedeclaration": { "limit": 8 @@ -105,10 +105,10 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 39011 + "limit": 39009 }, "reportUnknownParameterType": { - "limit": 19885 + "limit": 19883 }, "reportUnknownVariableType": { "limit": 30569 diff --git a/litellm/__init__.py b/litellm/__init__.py index e95b553c5d4..a520d284906 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -199,6 +199,7 @@ standard_logging_payload_excluded_fields: Optional[List[str]] = ( None # Fields to exclude from StandardLoggingPayload before callbacks receive it ) log_raw_request_response: bool = False +log_client_error_tracebacks: bool = False request_correlation_in_logs: bool = False redact_messages_in_exceptions: Optional[bool] = False redact_user_api_key_info: Optional[bool] = False @@ -1628,6 +1629,9 @@ if TYPE_CHECKING: AmazonMantleMessagesConfig as AmazonMantleMessagesConfig, ) from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig + from .llms.together_ai.chat.transformation import ( + TogetherAIChatConfig as TogetherAIChatConfig, + ) from .llms.nlp_cloud.chat.handler import NLPCloudConfig as NLPCloudConfig from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig as VertexGeminiConfig, @@ -1801,6 +1805,9 @@ if TYPE_CHECKING: from .llms.gemini.interactions.transformation import ( GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig, ) + from .llms.vertex_ai.interactions.transformation import ( + VertexAIInteractionsConfig as VertexAIInteractionsConfig, + ) from .llms.openai.chat.o_series_transformation import ( OpenAIOSeriesConfig as OpenAIOSeriesConfig, OpenAIOSeriesConfig as OpenAIO1Config, diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 89c72acc06d..1c833256598 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -177,6 +177,7 @@ LLM_CONFIG_NAMES: Final = ( "AmazonAnthropicClaudeMessagesConfig", "AmazonMantleMessagesConfig", "TogetherAIConfig", + "TogetherAIChatConfig", "NLPCloudConfig", "VertexGeminiConfig", "GoogleAIStudioGeminiConfig", @@ -242,6 +243,7 @@ LLM_CONFIG_NAMES: Final = ( "OpenRouterResponsesAPIConfig", "BedrockMantleResponsesAPIConfig", "GoogleAIStudioInteractionsConfig", + "VertexAIInteractionsConfig", "OpenAIOSeriesConfig", "AnthropicSkillsConfig", "BaseSkillsAPIConfig", @@ -740,6 +742,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "AmazonMantleMessagesConfig", ), "TogetherAIConfig": (".llms.together_ai.chat", "TogetherAIConfig"), + "TogetherAIChatConfig": ( + ".llms.together_ai.chat.transformation", + "TogetherAIChatConfig", + ), "NLPCloudConfig": (".llms.nlp_cloud.chat.handler", "NLPCloudConfig"), "VertexGeminiConfig": ( ".llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini", @@ -977,6 +983,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.gemini.interactions.transformation", "GoogleAIStudioInteractionsConfig", ), + "VertexAIInteractionsConfig": ( + ".llms.vertex_ai.interactions.transformation", + "VertexAIInteractionsConfig", + ), "OpenAIOSeriesConfig": ( ".llms.openai.chat.o_series_transformation", "OpenAIOSeriesConfig", diff --git a/litellm/caching/redis_cluster_node_isolation.py b/litellm/caching/redis_cluster_node_isolation.py index 8b0c120e80c..ae8c78709d9 100644 --- a/litellm/caching/redis_cluster_node_isolation.py +++ b/litellm/caching/redis_cluster_node_isolation.py @@ -18,6 +18,14 @@ already does when one of its pooled connections errors), leaving every other nod connections untouched. Every other branch (MOVED, ASK, CLUSTERDOWN, slot-not-covered, retry-exhaustion) is unchanged from upstream, since those already carry real evidence the topology changed. + +redis-py 8.x fixed this upstream with gentler machinery than this override's +``node.disconnect()`` (which also kills connections other coroutines are mid-operation +on, so one timeout cascades into a reconnect storm and, with TLS, a fresh handshake per +killed connection): it marks in-use connections for reconnect only after their current +operation completes, disconnects only the idle pooled ones, and defers reinitialization +to the outer retry loop. When the installed ``ClusterNode`` has that per-connection +recovery API, the factory returns the base ``RedisCluster`` unmodified. """ import asyncio @@ -72,8 +80,16 @@ class _ClusterAttrs(Protocol): _VERIFIED_REDIS_VERSIONS: Final = frozenset({"5.3.1"}) -def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]: - """Builds the ``RedisCluster`` subclass with the per-node isolation fix. +def get_litellm_async_redis_cluster_class( + cluster_node_class: type | None = None, +) -> type["_AsyncRedisClusterType"]: + """Returns the base ``RedisCluster`` when the installed redis-py already recovers a + node-level connection error per-connection (8.x+), else builds the ``RedisCluster`` + subclass with the per-node isolation fix for older versions whose upstream branch + tears down the whole cluster client. + + ``cluster_node_class`` exists for dependency injection in tests; production callers + leave it unset and the installed ``ClusterNode`` is used. Imported lazily because this module is reachable from a base ``import litellm`` while redis is not a base dependency. Cheap to call repeatedly: the underlying redis @@ -81,7 +97,10 @@ def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]: """ import redis from redis.asyncio.cluster import ( - RedisCluster as _BaseAsyncRedisCluster, # pyright: ignore[reportUnknownVariableType] # redis-py ships no resolvable stub for this class under the repo's current (stale) types-redis pin + ClusterNode as _AsyncClusterNode, # pyright: ignore[reportUnknownVariableType] # redis-py ships no resolvable stub for this class under the repo's current (stale) types-redis pin + ) + from redis.asyncio.cluster import ( + RedisCluster as _BaseAsyncRedisCluster, # pyright: ignore[reportUnknownVariableType] # same stale-stub gap as the import above ) from redis.cluster import get_node_name from redis.commands import READ_COMMANDS @@ -98,6 +117,15 @@ def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]: from redis.exceptions import ConnectionError as _RedisConnectionError from redis.exceptions import TimeoutError as _RedisTimeoutError + node_class: Final = cluster_node_class if cluster_node_class is not None else _AsyncClusterNode + if hasattr(node_class, "update_active_connections_for_reconnect"): + verbose_logger.debug( + "redis-py %s recovers a node-level connection error per-connection upstream; " + "using the base RedisCluster without litellm's node-isolation override.", + redis.__version__, + ) + return _BaseAsyncRedisCluster + if redis.__version__ not in _VERIFIED_REDIS_VERSIONS: verbose_logger.warning( "redis-py %s is not in the set this cluster-teardown-storm fix was verified " diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index b94e91b3034..17815976b4a 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -5,7 +5,7 @@ Handler for transforming /chat/completions api requests to litellm.responses req import json import os from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast, get_args from openai.types.responses.custom_tool_param import CustomToolParam from openai.types.responses.response_input_param import ( @@ -35,6 +35,7 @@ from litellm.responses.sse_output_recovery import ( ) from litellm.responses.utils import normalize_responses_api_stream_options from litellm.types.llms.openai import ( + REASONING_EFFORT, ChatCompletionAnnotation, ChatCompletionReasoningItem, ChatCompletionToolCallChunk, @@ -1113,22 +1114,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true" ) - # If string is passed, map with optional summary based on flag/env var - if reasoning_effort == "none": - return Reasoning(effort="none", summary="detailed") if auto_summary_enabled else Reasoning(effort="none") - elif reasoning_effort == "high": - return Reasoning(effort="high", summary="detailed") if auto_summary_enabled else Reasoning(effort="high") - elif reasoning_effort == "xhigh": - return Reasoning(effort="xhigh", summary="detailed") if auto_summary_enabled else Reasoning(effort="xhigh") - elif reasoning_effort == "medium": + if reasoning_effort in get_args(REASONING_EFFORT): return ( - Reasoning(effort="medium", summary="detailed") if auto_summary_enabled else Reasoning(effort="medium") - ) - elif reasoning_effort == "low": - return Reasoning(effort="low", summary="detailed") if auto_summary_enabled else Reasoning(effort="low") - elif reasoning_effort == "minimal": - return ( - Reasoning(effort="minimal", summary="detailed") if auto_summary_enabled else Reasoning(effort="minimal") + Reasoning(effort=reasoning_effort, summary="detailed") + if auto_summary_enabled + else Reasoning(effort=reasoning_effort) ) return None diff --git a/litellm/constants.py b/litellm/constants.py index 0a1ada3bab2..765bbfe1e54 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -147,6 +147,7 @@ LITELLM_UI_ALLOW_HEADERS: Final = [ "x-litellm-adaptive-router-model", "x-litellm-applied-guardrails", "x-litellm-guardrail-scan-id", + "x-litellm-cache-key", ] # Gemini model-specific minimal thinking budget constants @@ -750,6 +751,7 @@ openai_compatible_endpoints: Final[list] = [ "api.groq.com/openai/v1", "https://integrate.api.nvidia.com/v1", "api.deepseek.com/v1", + "api.together.ai/v1", "api.together.xyz/v1", "app.empower.dev/api/v1", "https://api.friendli.ai/serverless/v1", diff --git a/litellm/interactions/utils.py b/litellm/interactions/utils.py index 8a1e8836894..3895a85061d 100644 --- a/litellm/interactions/utils.py +++ b/litellm/interactions/utils.py @@ -47,6 +47,13 @@ def get_provider_interactions_api_config( return GoogleAIStudioInteractionsConfig() + if provider in (LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value): + from litellm.llms.vertex_ai.interactions.transformation import ( + VertexAIInteractionsConfig, + ) + + return VertexAIInteractionsConfig() + return None diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index de1092bc02f..b71ef4c8e06 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -58,6 +58,26 @@ def safe_divide( return numerator / denominator +def is_expected_client_error(exception: BaseException | None) -> bool: + """ + True when the exception maps to an HTTP 4xx status. + + ProxyException stores the status on .code (as a str), HTTPException and + litellm exceptions on .status_code. + """ + if exception is None: + return False + code: Final[object] = getattr(exception, "code", None) + status_code: Final[object] = code if code is not None else getattr(exception, "status_code", None) + if status_code is None or isinstance(status_code, bool): + return False + try: + status: Final = int(str(status_code)) + except ValueError: + return False + return 400 <= status < 500 + + def coerce_token_limit(value: object) -> int | None: """ Coerce a max_input_tokens / max_output_tokens value to an int, treating a diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index e674fc37673..005e94ebe82 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -272,6 +272,14 @@ def get_llm_provider( elif endpoint == "api.deepseek.com/v1": custom_llm_provider = "deepseek" dynamic_api_key = get_secret_str("DEEPSEEK_API_KEY") + elif endpoint == "api.together.ai/v1" or endpoint == "api.together.xyz/v1": + custom_llm_provider = "together_ai" + dynamic_api_key = api_key or ( + get_secret_str("TOGETHER_API_KEY") + or get_secret_str("TOGETHER_AI_API_KEY") + or get_secret_str("TOGETHERAI_API_KEY") + or get_secret_str("TOGETHER_AI_TOKEN") + ) elif endpoint == "ollama.com": custom_llm_provider = "ollama" dynamic_api_key = get_secret_str("OLLAMA_API_KEY") @@ -707,7 +715,7 @@ def _get_openai_compatible_provider_info( dynamic_api_key, ) = litellm.ZAIChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "together_ai": - api_base = api_base or get_secret_str("TOGETHER_AI_API_BASE") or "https://api.together.xyz/v1" + api_base = api_base or get_secret_str("TOGETHER_AI_API_BASE") or "https://api.together.ai/v1" dynamic_api_key = api_key or ( get_secret_str("TOGETHER_API_KEY") or get_secret_str("TOGETHER_AI_API_KEY") diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index 72f36661f4c..7a16ffe4d85 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -172,7 +172,7 @@ def get_supported_openai_params( if request_type == "embeddings": return litellm.JinaAIEmbeddingConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "together_ai": - return litellm.TogetherAIConfig().get_supported_openai_params(model=model) + return litellm.TogetherAIChatConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "databricks": if request_type == "chat_completion": return litellm.DatabricksConfig().get_supported_openai_params(model=model) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5437ce52706..626af7530a4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -62,7 +62,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.deepeval.deepeval import DeepEvalLogger from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger -from litellm.litellm_core_utils.core_helpers import reconstruct_model_name +from litellm.litellm_core_utils.core_helpers import is_expected_client_error, reconstruct_model_name from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( cost_breakdown_with_guardrail, @@ -3124,6 +3124,13 @@ class Logging(LiteLLMLoggingBaseClass): if not hasattr(self, "model_call_details"): self.model_call_details = {} + if ( + self.model_call_details.get("log_event_type") == "failed_api_call" + and self.model_call_details.get("exception") is exception + and self.model_call_details.get("standard_logging_object") is not None + ): + return start_time, self.model_call_details["end_time"] + self.model_call_details["log_event_type"] = "failed_api_call" self.model_call_details["exception"] = exception self.model_call_details["traceback_exception"] = ( @@ -5455,9 +5462,10 @@ class StandardLoggingPayloadSetup: error_class: Final[str] = str(original_exception.__class__.__name__) if original_exception else "" _llm_provider_in_exception: Final = getattr(original_exception, "llm_provider", "") - # Get traceback information (first 100 lines) traceback_info = traceback_str or "" - if original_exception: + if original_exception and ( + litellm.log_client_error_tracebacks or not is_expected_client_error(original_exception) + ): tb: Final[TracebackType | None] = getattr(original_exception, "__traceback__", None) if tb: tb_lines: Final = traceback.format_tb(tb) @@ -5930,11 +5938,15 @@ def get_standard_logging_object_payload( response_model_name = final_response_obj.get("model") # For Azure Model Router, preserve the actual model in the top-level standard - # logging payload only when the user has opted in. + # logging payload. + from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + requested_model: Final = kwargs.get("model") - if ( - isinstance(requested_model, str) - and ("model_router" in requested_model.lower() or "model-router" in requested_model.lower()) + stamped_selected_model: Final = AzureFoundryModelInfo.get_model_router_selected_model(hidden_params) + if stamped_selected_model is not None: + model_name = stamped_selected_model + elif ( + AzureFoundryModelInfo.is_model_router_call(model=requested_model, hidden_params=hidden_params) and isinstance(response_model_name, str) and response_model_name ): diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 826a890eca9..86cfbf70255 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -643,49 +643,6 @@ def claude_2_1_pt( return prompt -### TOGETHER AI - - -def get_model_info(token, model): - try: - headers: Final = {"Authorization": f"Bearer {token}"} - client: Final = HTTPHandler(concurrent_limit=1) - response: Final = client.get("https://api.together.xyz/models/info", headers=headers) - if response.status_code == 200: - model_info: Final = response.json() - for m in model_info: - if m["name"].lower().strip() == model.strip(): - return m["config"].get("prompt_format", None), m["config"].get("chat_template", None) - return None, None - else: - return None, None - except Exception: # safely fail a prompt template request - return None, None - - -## OLD TOGETHER AI FLOW -# def format_prompt_togetherai(messages, prompt_format, chat_template): -# if prompt_format is None: -# return default_pt(messages) - -# human_prompt, assistant_prompt = prompt_format.split("{prompt}") - -# if chat_template is not None: -# prompt = hf_chat_template( -# model=None, messages=messages, chat_template=chat_template -# ) -# elif prompt_format is not None: -# prompt = custom_prompt( -# role_dict={}, -# messages=messages, -# initial_prompt_value=human_prompt, -# final_prompt_value=assistant_prompt, -# ) -# else: -# prompt = default_pt(messages) -# return prompt - - ### IBM Granite diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 23abca7d5f2..26380ad0af8 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1215,8 +1215,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if reasoning_effort is None or reasoning_effort == "none": return None if AnthropicConfig._is_adaptive_thinking_model(model, custom_llm_provider): + # without display, Anthropic defaults adaptive thinking to + # display="omitted" and returns a blank thinking block return AnthropicThinkingParam( type="adaptive", + display="summarized", ) elif reasoning_effort == "low": return AnthropicThinkingParam( @@ -2144,7 +2147,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) @staticmethod - def _thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None: + def thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None: details: Final = usage_object.get("output_tokens_details") if not isinstance(details, Mapping): return None @@ -2176,7 +2179,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): reported_thinking_tokens: Final = ( iteration_thinking_tokens if iteration_thinking_tokens is not None - else self._thinking_tokens_from_usage(usage_object) + else self.thinking_tokens_from_usage(usage_object) ) if reported_thinking_tokens is not None: capped_reported: Final = min(max(0, reported_thinking_tokens), completion_tokens) @@ -2199,7 +2202,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _sum_iteration_thinking_tokens(self, iterations: Sequence[object]) -> int | None: per_iteration: Final = tuple( - self._thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None + self.thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None for iteration in iterations ) reported: Final = tuple(tokens for tokens in per_iteration if tokens is not None) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 9461e40cf2e..53cc3464761 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -38,6 +38,21 @@ DROP_DISABLED_THINKING_WARNING: Final = ( "thinking blocks, and those thinking tokens are billed as output tokens." ) +# Anthropic error `type` (both the JSON error body and SSE `event: error` +# payloads use this field) mapped to the HTTP status code it corresponds to. +ANTHROPIC_ERROR_STATUS_CODE_MAP: Final = MappingProxyType( + { + "invalid_request_error": 400, + "authentication_error": 401, + "permission_error": 403, + "not_found_error": 404, + "rate_limit_error": 429, + "api_error": 500, + "overloaded_error": 503, + "timeout_error": 504, + } +) + _BEDROCK_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$") _INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$") _DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$") diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 7c89da81fe6..109017bda27 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -434,7 +434,7 @@ class LiteLLMAnthropicMessagesAdapter: content_items = list(content.get("content", [])) # Single-item text keeps the backward-compatible string format; a single - # image becomes a structured image_url part + # image or document becomes a structured image_url part if len(content_items) == 1: c = content_items[0] if isinstance(c, str): @@ -454,7 +454,7 @@ class LiteLLMAnthropicMessagesAdapter: ) self._add_cache_control_if_applicable(content, tool_result, model) tool_message_list.append(tool_result) - elif c.get("type") == "image": + elif c.get("type") in ("image", "document"): image_part = self._tool_result_image_part(c.get("source")) tool_result = ChatCompletionToolMessage( role="tool", @@ -482,7 +482,7 @@ class LiteLLMAnthropicMessagesAdapter: text=c.get("text", ""), ) ) - elif c.get("type") == "image": + elif c.get("type") in ("image", "document"): image_part = self._tool_result_image_part(c.get("source")) if image_part: combined_content_parts.append(image_part) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 922769dbbfd..0a12bc3135f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -1,6 +1,6 @@ import asyncio import json -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from datetime import datetime from typing import Any, Final, Protocol, runtime_checkable @@ -11,9 +11,11 @@ from typing_extensions import TypedDict from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.anthropic.common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) +from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType from litellm.types.utils import GenericStreamingChunk, ModelResponseStream @@ -33,26 +35,239 @@ def _is_message_stop_chunk(chunk: object) -> bool: return False -def _is_provider_error_chunk(chunk: object) -> bool: +def is_anthropic_ping_chunk(chunk: object) -> bool: + """ + Whether a chunk is a pure ``ping`` keepalive frame. It carries no content + and can recur indefinitely on a slow-starting or idle connection, so a + mid-stream fallback wrapper drops it outright while still deciding + whether to commit to the primary stream, rather than buffering it. + + A physical transport chunk that coalesces a ping with any other SSE + event (``message_start``, ``content_block_delta``, ``event: error``, ...) + is NOT a pure ping - dropping it whole would discard those events - so + only a chunk whose every ``event:`` line is ``event: ping`` qualifies. + """ if isinstance(chunk, dict): - return chunk.get("type") == "error" + return chunk.get("type") == "ping" if isinstance(chunk, (bytes, bytearray)): - return any(line == b"event: error" for line in chunk.splitlines()) + event_lines: Final = tuple(line for line in chunk.splitlines() if line.startswith(b"event:")) + return bool(event_lines) and all(line == b"event: ping" for line in event_lines) return False +def is_anthropic_content_delta_chunk(chunk: object) -> bool: + """ + Whether a chunk carries actual assistant-generated output (a + ``content_block_delta`` frame), as opposed to a lifecycle/bookkeeping + frame (``message_start``, ``content_block_start``/``stop``, + ``message_delta``, ``message_stop``, ``ping``) that carries nothing + worth preserving before an invisible mid-stream fallback retry. + """ + if isinstance(chunk, dict): + return chunk.get("type") == "content_block_delta" + if isinstance(chunk, (bytes, bytearray)): + return any(line == b"event: content_block_delta" for line in chunk.splitlines()) + return False + + +def _decoded_sse_data_line(line: bytes) -> object | None: + if not line.startswith(b"data:"): + return None + try: + return json.loads(line[len(b"data:") :].strip()) + except (ValueError, TypeError): + return None + + +def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None: + if isinstance(chunk, dict): + return chunk if chunk.get("type") == "error" else None + if isinstance(chunk, (bytes, bytearray)): + decoded_lines: Final = (_decoded_sse_data_line(line) for line in chunk.splitlines()) + return next( + ( + candidate + for candidate in decoded_lines + if isinstance(candidate, dict) and candidate.get("type") == "error" + ), + None, + ) + return None + + +def _anthropic_error_body(chunk: object) -> Mapping[str, object] | None: + """Return the ``error`` object of an Anthropic SSE ``event: error`` chunk, or None.""" + payload: Final = _anthropic_error_event_payload(chunk) + error_body: Final = payload.get("error") if payload is not None else None + return error_body if isinstance(error_body, dict) else None + + +def _is_provider_error_chunk(chunk: object) -> bool: + return _anthropic_error_body(chunk) is not None + + +def parse_anthropic_error_event(chunk: object) -> tuple[str, str, int] | None: + """ + Extract ``(error_type, message, http_status_code)`` from an Anthropic SSE + ``event: error`` chunk (raw bytes or an already-decoded dict), or None if + ``chunk`` is not an error event. + + The status code is looked up via ANTHROPIC_ERROR_STATUS_CODE_MAP, + defaulting to 500 for an error ``type`` Anthropic hasn't documented yet. + """ + error_body: Final = _anthropic_error_body(chunk) + if error_body is None: + return None + error_type: Final = error_body.get("type") + if not isinstance(error_type, str): + return None + message: Final = error_body.get("message") + return ( + error_type, + message if isinstance(message, str) else error_type, + ANTHROPIC_ERROR_STATUS_CODE_MAP.get(error_type, 500), + ) + + def _is_terminal_stream_chunk(chunk: object) -> bool: return _is_message_stop_chunk(chunk) or _is_provider_error_chunk(chunk) +def _sse_event(event_type: str, payload: Mapping[str, object]) -> bytes: + return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode() + + def _incomplete_stream_error_sse_event() -> bytes: - payload: Final = json.dumps( - { - "type": "error", - "error": {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE}, - } + return _sse_event( # mutable-ok: one-shot JSON payload, never mutated after construction + "error", + {"type": "error", "error": {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE}}, + ) + + +def _anthropic_content_block_start_and_deltas( + block: Mapping[str, object], +) -> tuple[Mapping[str, object], tuple[Mapping[str, object], ...]]: + """ + ``(content_block_start.content_block, content_block_delta.delta events)`` + for one Anthropic response content block. A thinking block emits both a + thinking_delta and a trailing signature_delta - a real Anthropic stream + does the same, and dropping the signature makes any replay of that + assistant message (a follow-up turn, a tool-use continuation) fail + Anthropic's thinking-signature verification. redacted_thinking has no + delta at all - it is sent complete in content_block_start. + """ + match block.get("type"): + case "tool_use": + return ( + { # mutable-ok: one-shot payload + "id": block.get("id"), + "name": block.get("name"), + "input": {}, # mutable-ok: one-shot payload + "type": "tool_use", + }, + ( + { # mutable-ok: one-shot payload + "partial_json": json.dumps(block.get("input") or {}), # mutable-ok: one-shot payload + "type": "input_json_delta", + }, + ), + ) + case "thinking": + signature: Final = block.get("signature") + signature_deltas: Final = ( + ({"signature": signature, "type": "signature_delta"},) # mutable-ok: one-shot payload + if isinstance(signature, str) and signature + else () + ) + return ( + {"thinking": "", "signature": "", "type": "thinking"}, # mutable-ok: one-shot payload + ( + {"thinking": block.get("thinking") or "", "type": "thinking_delta"}, # mutable-ok: one-shot payload + *signature_deltas, + ), + ) + case "redacted_thinking": + return ({"type": "redacted_thinking", "data": block.get("data")}, ()) # mutable-ok: one-shot JSON payload + case _: + return ( + {"type": "text", "text": ""}, # mutable-ok: one-shot JSON payload + ({"type": "text_delta", "text": block.get("text") or ""},), # mutable-ok: one-shot JSON payload + ) + + +def anthropic_messages_response_as_sse_events(response: AnthropicMessagesResponse) -> tuple[bytes, ...]: + """ + Render a complete (non-streaming) AnthropicMessagesResponse as the SSE + event sequence a real streaming request would have produced. + + A mid-stream fallback can resolve to a non-streaming response even + though the client asked to stream (e.g. an agentic tool-use loop that + intercepts and returns a complete message) - yielding that dict directly + into a `/v1/messages` SSE byte stream would produce a malformed + response, so it's synthesized into the message_start/content_block_*/ + message_delta/message_stop lifecycle a real stream would have sent. + """ + content_blocks: Final = response.get("content") or () + content_events: Final = ( + event for index, block in enumerate(content_blocks) for event in _anthropic_content_block_events(index, block) + ) + # A real message_start always carries a null stop_reason/stop_sequence and + # a zero output_tokens - those are only known once generation finishes, so + # copying the completed response's final values here would let a client + # treat the message as already finished, or double-count output tokens. + message_start_usage: Final = { # mutable-ok: one-shot JSON payload + **(response.get("usage") or {}), + "output_tokens": 0, + } + message_start_payload: Final = { # mutable-ok: one-shot JSON payload, never mutated after construction + "type": "message_start", + "message": { # mutable-ok: one-shot JSON payload + **response, + "content": [], # mutable-ok: one-shot JSON payload + "stop_reason": None, + "stop_sequence": None, + "usage": message_start_usage, + }, + } + message_delta_payload: Final = { # mutable-ok: one-shot JSON payload, never mutated after construction + "type": "message_delta", + "delta": { # mutable-ok: one-shot JSON payload + "stop_reason": response.get("stop_reason"), + "stop_sequence": response.get("stop_sequence"), + }, + "usage": response.get("usage") or {}, # mutable-ok: one-shot JSON payload + } + return ( + _sse_event("message_start", message_start_payload), + *content_events, + _sse_event("message_delta", message_delta_payload), + _sse_event("message_stop", {"type": "message_stop"}), # mutable-ok: one-shot JSON payload + ) + + +def _anthropic_content_block_events(index: int, block: Mapping[str, object]) -> tuple[bytes, ...]: + start_block, deltas = _anthropic_content_block_start_and_deltas(block) + start_payload: Final = { # mutable-ok: one-shot payload + "type": "content_block_start", + "index": index, + "content_block": start_block, + } + stop_payload: Final = { # mutable-ok: one-shot payload + "type": "content_block_stop", + "index": index, + } + delta_events: Final = tuple( + _sse_event( + "content_block_delta", + {"type": "content_block_delta", "index": index, "delta": delta}, # mutable-ok: one-shot payload + ) + for delta in deltas + ) + return ( + _sse_event("content_block_start", start_payload), + *delta_events, + _sse_event("content_block_stop", stop_payload), ) - return f"event: error\ndata: {payload}\n\n".encode() class AnthropicMessagesStreamHiddenParams(TypedDict): diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 6d47d0de19f..492e0050cfe 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -87,6 +87,51 @@ class LiteLLMAnthropicToResponsesAPIAdapter: return source.get("url") return None + @staticmethod + def _translate_anthropic_document_block_to_file_part( + block: Mapping[str, object], + ) -> dict[str, str] | None: # mutable-ok: API message payload + """Convert an Anthropic document block to a Responses input_file part.""" + raw_source: Final = block.get("source") + if not isinstance(raw_source, Mapping): + return None + source: Final = cast(Mapping[str, object], raw_source) # cast-ok: untrusted client payload + source_type: Final = source.get("type") + if source_type == "base64": + data: Final = source.get("data") + if not isinstance(data, str) or not data: + return None + raw_media_type: Final = source.get("media_type") + media_type: Final = ( + raw_media_type if isinstance(raw_media_type, str) and raw_media_type else "application/pdf" + ) + raw_title: Final = block.get("title") + filename: Final = raw_title if isinstance(raw_title, str) and raw_title else "document.pdf" + return { # mutable-ok: API message payload + "type": "input_file", + "filename": filename, + "file_data": f"data:{media_type};base64,{data}", + } + if source_type == "url": + url: Final = source.get("url") + if not isinstance(url, str) or not url: + return None + return {"type": "input_file", "file_url": url} # mutable-ok: API message payload + return None + + @staticmethod + def _tool_result_output_value( + output_text: str, + file_parts: tuple[dict[str, str], ...], # mutable-ok: json content parts + ) -> str | list[dict[str, str]]: # mutable-ok: API message payload + """Plain string output, or a part list when document file parts are present.""" + if not file_parts: + return output_text + text_parts: Final = ( + [{"type": "input_text", "text": output_text}] if output_text else [] # mutable-ok: API message payload + ) + return [*text_parts, *file_parts] # mutable-ok: API message payload + @staticmethod def _translate_midturn_system_content_to_responses( content: str | Iterable[AnthropicSystemMessageContent], @@ -169,6 +214,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: system text -> message(role=system, input_text) user text -> message(role=user, input_text) user image -> message(role=user, input_image) + user document -> message(role=user, input_file) user tool_result -> function_call_output assistant text -> message(role=assistant, output_text) assistant thinking -> reasoning @@ -223,9 +269,25 @@ class LiteLLMAnthropicToResponsesAPIAdapter: {"type": "input_image", "image_url": url}, block.get("prompt_cache_breakpoint") ) ) + elif btype == "document": + file_part = self._translate_anthropic_document_block_to_file_part(block) + if file_part: + user_parts.append( + with_prompt_cache_breakpoint(file_part, block.get("prompt_cache_breakpoint")) + ) elif btype == "tool_result": tool_use_id = block.get("tool_use_id", "") inner = block.get("content") + document_candidates = ( + tuple( + self._translate_anthropic_document_block_to_file_part(c) + for c in inner + if isinstance(c, dict) and c.get("type") == "document" + ) + if isinstance(inner, list) + else () + ) + tool_file_parts = tuple(part for part in document_candidates if part is not None) if inner is None: output_text = "" elif isinstance(inner, str): @@ -258,7 +320,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: { "type": "function_call_output", "call_id": tool_use_id, - "output": output_text, + "output": self._tool_result_output_value(output_text, tool_file_parts), } ) if tool_image_parts: diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index 0c62418708f..5fdf2ceff7f 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -22,19 +22,7 @@ from litellm.types.llms.openai import ( from litellm.types.utils import CallTypes, LlmProviders, ModelResponse from ..chat.transformation import AnthropicConfig -from ..common_utils import AnthropicModelInfo - -# Map Anthropic error types to HTTP status codes -ANTHROPIC_ERROR_STATUS_CODE_MAP: Final = { - "invalid_request_error": 400, - "authentication_error": 401, - "permission_error": 403, - "not_found_error": 404, - "rate_limit_error": 429, - "api_error": 500, - "overloaded_error": 503, - "timeout_error": 504, -} +from ..common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP, AnthropicModelInfo class AnthropicFilesHandler: diff --git a/litellm/llms/azure_ai/azure_model_router/transformation.py b/litellm/llms/azure_ai/azure_model_router/transformation.py index 1a924088390..61cbc213b11 100644 --- a/litellm/llms/azure_ai/azure_model_router/transformation.py +++ b/litellm/llms/azure_ai/azure_model_router/transformation.py @@ -65,15 +65,24 @@ class AzureModelRouterConfig(AzureAIStudioConfig): Extracts the actual model used from the Azure response (e.g., gpt-5-nano-2025-08-07) and returns it with the azure_ai/ prefix for proper display and cost tracking. + + Also stamps that model onto ``_hidden_params`` so downstream consumers (spend logs, + response restamping) can read it instead of guessing the route from the model string. """ - from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + AzureFoundryModelInfo, + ) + from litellm.router_utils.add_retry_fallback_headers import ( + get_hidden_params_dict, + ) # Get base model for the parent call (strips routing prefixes for API compatibility) base_model: Final[str] = AzureFoundryModelInfo.get_base_model(model) # Call parent transform_response first - this will extract the actual model # from the raw response (e.g., "gpt-5-nano-2025-08-07") - model_response = super().transform_response( + transformed_response: Final = super().transform_response( model=base_model, raw_response=raw_response, model_response=model_response, @@ -86,7 +95,15 @@ class AzureModelRouterConfig(AzureAIStudioConfig): api_key=api_key, json_mode=json_mode, ) - return model_response + selected_model: Final = transformed_response.model + if selected_model: + # Rebuilt rather than mutated in place: ModelResponseBase declares _hidden_params as a + # class-level dict, so an in-place write can bleed into unrelated responses. + transformed_response._hidden_params = { # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params setter # mutable-ok: ModelResponse requires _hidden_params to be a plain dict + **get_hidden_params_dict(transformed_response), + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model, + } + return transformed_response def calculate_additional_costs(self, model: str, prompt_tokens: int, completion_tokens: int) -> dict | None: """ diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index d09055d7671..26a90157455 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -51,6 +51,9 @@ def get_azure_ai_auth_headers( ) +AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model" + + class AzureFoundryModelInfo(BaseLLMModelInfo): """Model info for Azure AI / Azure Foundry models.""" @@ -82,6 +85,41 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): return "model_router" return "default" + @staticmethod + def get_model_router_selected_model(hidden_params: Mapping[str, object] | None) -> str | None: + """The model Azure Model Router actually served, stamped by ``AzureModelRouterConfig``. + + Reading this beats re-deriving the route from a model string: the stamp is set on the + code path that was actually taken, so it holds no matter what the caller named the model. + """ + if not hidden_params: + return None + selected: Final = hidden_params.get(AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY) + if isinstance(selected, str) and selected: + return selected + return None + + @staticmethod + def is_model_router_call( + model: str | None = None, + hidden_params: Mapping[str, object] | None = None, + ) -> bool: + """Whether a request went down the Azure Model Router route. + + Prefers the response stamp, then the deployment's litellm model path, and only then the + caller-supplied name. The last two go through ``get_azure_ai_route`` so the model-router + name heuristic lives in exactly one place. + """ + if AzureFoundryModelInfo.get_model_router_selected_model(hidden_params) is not None: + return True + deployment_model: Final = ( + hidden_params.get("litellm_model_name") or hidden_params.get("model") if hidden_params is not None else None + ) + return any( + isinstance(candidate, str) and AzureFoundryModelInfo.get_azure_ai_route(candidate) == "model_router" + for candidate in (deployment_model, model) + ) + @staticmethod def get_api_base(api_base: str | None = None) -> str | None: return api_base or litellm.api_base or get_secret_str("AZURE_AI_API_BASE") diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index b437e25d24b..767677cbcbf 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1617,6 +1617,8 @@ class AmazonConverseConfig(BaseConfig): } if additional_request_params: data["additionalModelRequestFields"] = additional_request_params + if "thinking" in additional_request_params: + data["additionalModelResponseFieldPaths"] = ("/usage/output_tokens_details",) if system_content_blocks: data["system"] = system_content_blocks @@ -1801,6 +1803,17 @@ class AmazonConverseConfig(BaseConfig): thinking_blocks_list.append(_redacted_block) return thinking_blocks_list + @staticmethod + def thinking_tokens_from_additional_fields(additional_fields: object) -> int | None: + """Converse omits thinking tokens from its usage block; they only arrive under + ``additionalModelResponseFields`` when ``/usage/output_tokens_details`` is requested.""" + if not isinstance(additional_fields, Mapping): + return None + usage: Final = additional_fields.get("usage") + if not isinstance(usage, Mapping): + return None + return AnthropicConfig.thinking_tokens_from_usage(usage) + @staticmethod def is_converse_usage_shape(usage_object: Mapping[str, object]) -> bool: """Converse-family models report camelCase token counts, not Anthropic's snake_case.""" @@ -1842,6 +1855,7 @@ class AmazonConverseConfig(BaseConfig): usage: ConverseTokenUsageBlock, reasoning_content: str | None = None, thinking_ran: bool = False, + provider_reasoning_tokens: int | None = None, ) -> Usage: input_tokens = usage["inputTokens"] output_tokens: Final = usage["outputTokens"] @@ -1862,9 +1876,14 @@ class AmazonConverseConfig(BaseConfig): cache_creation_tokens=cache_creation_input_tokens, text_tokens=raw_input_tokens, ) - reasoning_tokens: Final = ( + estimated_reasoning_tokens: Final = ( token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 ) + reasoning_tokens: Final = ( + min(max(0, provider_reasoning_tokens), output_tokens) + if provider_reasoning_tokens is not None + else estimated_reasoning_tokens + ) completion_tokens_details: Final = ( CompletionTokensDetailsWrapper( reasoning_tokens=reasoning_tokens, @@ -2272,6 +2291,9 @@ class AmazonConverseConfig(BaseConfig): completion_response["usage"], reasoning_content=chat_completion_message.get("reasoning_content"), thinking_ran=reasoningContentBlocks is not None, + provider_reasoning_tokens=self.thinking_tokens_from_additional_fields( + completion_response.get("additionalModelResponseFields") + ), ) ## HANDLE TOOL CALLS diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 35e82f1961a..fc34e403beb 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -331,6 +331,7 @@ class AWSEventStreamDecoder: self.json_mode = json_mode self._current_tool_name: str | None = None self._thinking_ran = False + self._provider_reasoning_tokens: int | None = None def check_empty_tool_call_args(self) -> bool: """ @@ -559,10 +560,14 @@ class AWSEventStreamDecoder: tool_use = self._handle_converse_stop_event(content_block_index) elif "stopReason" in chunk_data: finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop")) + self._provider_reasoning_tokens = AmazonConverseConfig.thinking_tokens_from_additional_fields( + chunk_data.get("additionalModelResponseFields") + ) elif "usage" in chunk_data: usage = converse_config.transform_usage( chunk_data.get("usage", {}), thinking_ran=self._thinking_ran, + provider_reasoning_tokens=self._provider_reasoning_tokens, ) if thinking_blocks: self._thinking_ran = True diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index 0ce2e6f60d3..d0a3c37ffb3 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -1,4 +1,5 @@ import json +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, cast from httpx import Response @@ -93,6 +94,9 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD endpoint_url, ) + def get_bedrock_bearer_token(self, litellm_params: Mapping[str, object]) -> str | None: + return None + def sign_request( self, headers: dict, @@ -109,6 +113,7 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD request_data=request_data or {}, api_base=api_base, model=model, + api_key=self.get_bedrock_bearer_token(optional_params), ) def logging_non_streaming_response( diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index cb0473887ea..4860c99268e 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -29,6 +29,7 @@ class BedrockRerankHandler(BaseAWSLLM): async def arerank( self, prepared_request: BedrockPreparedRequest, + logging_obj: LitellmLogging, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ): @@ -40,6 +41,7 @@ class BedrockRerankHandler(BaseAWSLLM): headers=dict(prepared_request["prepped"].headers), data=prepared_request["body"], timeout=timeout, + logging_obj=logging_obj, ) response.raise_for_status() except httpx.HTTPStatusError as err: @@ -98,6 +100,7 @@ class BedrockRerankHandler(BaseAWSLLM): if _is_async: return self.arerank( prepared_request, + logging_obj=logging_obj, timeout=timeout, client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None, ) diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py index 889361cd808..d877fbb4e09 100644 --- a/litellm/llms/bedrock_mantle/common_utils.py +++ b/litellm/llms/bedrock_mantle/common_utils.py @@ -13,6 +13,7 @@ global state. """ import re +from collections.abc import Mapping from typing import Final from botocore.exceptions import ( @@ -31,30 +32,39 @@ BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1" MANTLE_HOST_RE: Final = re.compile(r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE) +def resolve_mantle_bearer_token(api_key: str | None) -> str | None: + return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + + +def resolve_mantle_region(params: Mapping[str, object]) -> str: + region: Final = params.get("aws_region_name") + if isinstance(region, str) and region: + BaseAWSLLM._validate_aws_region_name(region) + return region + api_base: Final = params.get("api_base") + base: Final = (api_base if isinstance(api_base, str) else None) or get_secret_str("BEDROCK_MANTLE_API_BASE") + if base: + match: Final = MANTLE_HOST_RE.match(base.rstrip("/")) + if match: + return match.group(1) + return ( + get_secret_str("BEDROCK_MANTLE_REGION") + or get_secret_str("AWS_REGION_NAME") + or get_secret_str("AWS_REGION") + or BEDROCK_MANTLE_DEFAULT_REGION + ) + + class BedrockMantleAuthMixin: _aws_signer: BaseAWSLLM @staticmethod def _resolve_bearer_token(api_key: str | None) -> str | None: - return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + return resolve_mantle_bearer_token(api_key) @staticmethod def _resolve_region(params: dict) -> str: - region: Final = params.get("aws_region_name") - if region: - BaseAWSLLM._validate_aws_region_name(region) - return region - base: Final = params.get("api_base") or get_secret_str("BEDROCK_MANTLE_API_BASE") - if base: - match: Final = MANTLE_HOST_RE.match(base.rstrip("/")) - if match: - return match.group(1) - return ( - get_secret_str("BEDROCK_MANTLE_REGION") - or get_secret_str("AWS_REGION_NAME") - or get_secret_str("AWS_REGION") - or BEDROCK_MANTLE_DEFAULT_REGION - ) + return resolve_mantle_region(params) def sign_request( self, diff --git a/litellm/llms/bedrock_mantle/passthrough/transformation.py b/litellm/llms/bedrock_mantle/passthrough/transformation.py new file mode 100644 index 00000000000..e6b831efa57 --- /dev/null +++ b/litellm/llms/bedrock_mantle/passthrough/transformation.py @@ -0,0 +1,71 @@ +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final, Literal, Optional + +from httpx import Response + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig +from litellm.llms.bedrock_mantle.common_utils import ( + MANTLE_HOST_RE, + resolve_mantle_bearer_token, + resolve_mantle_region, +) +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from litellm.types.utils import CostResponseTypes + + +class BedrockMantlePassthroughConfig(BedrockPassthroughConfig): + """Native Bedrock runtime passthrough (InvokeModel, Converse) for deployments declared as bedrock_mantle. + + The Mantle host only serves the OpenAI-compatible surface, so a Mantle api_base lends its region and the + request itself goes to bedrock-runtime, signed with the deployment's Bearer token or SigV4 credentials. + """ + + def _get_aws_region_name( + self, + optional_params: Mapping[str, object], + model: str | None = None, + model_id: str | None = None, + ) -> str: + return resolve_mantle_region(optional_params) + + def get_runtime_endpoint( + self, + api_base: str | None, + aws_bedrock_runtime_endpoint: str | None, + aws_region_name: str, + endpoint_type: Literal["runtime", "agent", "agentcore"] | None = "runtime", + ) -> tuple[str, str]: + is_mantle_host: Final = api_base is not None and MANTLE_HOST_RE.match(api_base.rstrip("/")) is not None + return super().get_runtime_endpoint( + api_base=None if is_mantle_host else api_base, + aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, + aws_region_name=aws_region_name, + endpoint_type=endpoint_type, + ) + + def get_bedrock_bearer_token(self, litellm_params: Mapping[str, object]) -> str | None: + api_key: Final = litellm_params.get("api_key") + return resolve_mantle_bearer_token(api_key if isinstance(api_key, str) else None) + + def logging_non_streaming_response( + self, + model: str, + custom_llm_provider: str, + httpx_response: Response, + request_data: dict, # mutable-ok: mirrors the inherited BedrockPassthroughConfig signature + logging_obj: Logging, + endpoint: str, + ) -> Optional["CostResponseTypes"]: + is_converse: Final = "invoke" not in endpoint and "converse" in endpoint + shape_provider: Final = LlmProviders.BEDROCK.value if is_converse else custom_llm_provider + return super().logging_non_streaming_response( + model=model, + custom_llm_provider=shape_provider, + httpx_response=httpx_response, + request_data=request_data, + logging_obj=logging_obj, + endpoint=endpoint, + ) diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 92da5835b2d..3e5dd4ff87d 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -15,8 +15,12 @@ role / access key / profile / web identity), signed via the shared BaseAWSLLM._sign_request after the request body is finalized. """ +import json +from collections.abc import Mapping from typing import Any, Final +from typing_extensions import ReadOnly, TypedDict + import litellm from litellm._logging import verbose_logger from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -50,6 +54,33 @@ _BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"}) _CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools" +_CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE: Final = "agent_message" +_CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE: Final = "context_compaction" +_CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: Final = "local_shell_call" + + +class _RewrittenOutputTextBlock(TypedDict): + type: ReadOnly[str] + text: ReadOnly[str] + + +class _RewrittenAssistantMessageItem(TypedDict): + type: ReadOnly[str] + role: ReadOnly[str] + content: ReadOnly[tuple[_RewrittenOutputTextBlock, ...]] + + +class _RewrittenCompactionItem(TypedDict): + type: ReadOnly[str] + encrypted_content: ReadOnly[str] + + +class _RewrittenFunctionCallItem(TypedDict): + type: ReadOnly[str] + call_id: ReadOnly[str] + name: ReadOnly[str] + arguments: ReadOnly[str] + class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPIConfig): def __init__( @@ -155,6 +186,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI headers: dict, ) -> dict: remaining_input, hoisted_tools = self._hoist_codex_additional_tools(input) + normalized_input: Final = self._normalize_codex_input_items(remaining_input) request_params: Final = ( { **response_api_optional_request_params, @@ -168,7 +200,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI ) return super().transform_responses_api_request( model=model, - input=remaining_input, + input=normalized_input, response_api_optional_request_params=request_params, litellm_params=litellm_params, headers=headers, @@ -210,6 +242,91 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI ) return remaining_input, cls._filter_unsupported_tools(hoisted_tools) + @staticmethod + def _agent_message_text(item: "Mapping[str, Any]") -> str: + content: Final = item.get("content") + if not isinstance(content, list): + return "" + return "".join( + str(block.get("text") or block.get("encrypted_content") or "") + for block in content + if isinstance(block, dict) + ) + + @classmethod + def _normalize_agent_message_item(cls, item: "Mapping[str, Any]") -> "_RewrittenAssistantMessageItem | None": + text: Final = cls._agent_message_text(item) + if not text: + return None + rewritten: Final[_RewrittenAssistantMessageItem] = { + "type": "message", + "role": "assistant", + "content": ({"type": "output_text", "text": text},), + } + return rewritten + + @staticmethod + def _normalize_context_compaction_item(item: "Mapping[str, Any]") -> "_RewrittenCompactionItem | None": + encrypted_content: Final = item.get("encrypted_content") + if not isinstance(encrypted_content, str) or not encrypted_content: + return None + rewritten: Final[_RewrittenCompactionItem] = {"type": "compaction", "encrypted_content": encrypted_content} + return rewritten + + @staticmethod + def _normalize_local_shell_call_item(item: "Mapping[str, Any]") -> "_RewrittenFunctionCallItem | None": + call_id: Final = item.get("call_id") + if not isinstance(call_id, str) or not call_id: + return None + action: Final = item.get("action") + rewritten: Final[_RewrittenFunctionCallItem] = { + "type": "function_call", + "call_id": call_id, + "name": "local_shell", + "arguments": json.dumps(action) if isinstance(action, dict) else "{}", + } + return rewritten + + @classmethod + def _normalize_codex_input_item(cls, item: object) -> "tuple[object, str | None]": + """Returns (normalized item or None to drop it, original type when rewritten).""" + if not isinstance(item, dict): + return item, None + item_type: Final = item.get("type") + if item_type == _CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE: + return cls._normalize_agent_message_item(item), item_type + if item_type == _CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE: + return cls._normalize_context_compaction_item(item), item_type + if item_type == _CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: + return cls._normalize_local_shell_call_item(item), item_type + return item, None + + @classmethod + def _normalize_codex_input_items( + cls, + input: "str | ResponseInputParam", + ) -> "str | ResponseInputParam": + """Rewrite Codex history item types Mantle rejects with 400 "Invalid + 'input': value did not match any expected variant" into supported + equivalents. `agent_message` (Codex multi-agent traffic; its + encrypted_content slot carries the plaintext payload when the model + never issued encrypted args) becomes an assistant message, + `context_compaction` becomes the `compaction` spelling Mantle accepts, + and `local_shell_call` becomes the function_call its recorded + function_call_output already pairs with. + """ + if not isinstance(input, list): + return input + normalized: Final = tuple(cls._normalize_codex_input_item(item) for item in input) + rewritten_types: Final = sorted(frozenset(item_type for _, item_type in normalized if item_type is not None)) + if rewritten_types: + verbose_logger.warning( + "Bedrock Mantle Responses API: rewrote Codex input item type(s) %s that Mantle rejects.", + rewritten_types, + ) + kept: Final = [item for item, _ in normalized if item is not None] # mutable-ok: ResponseInputParam is a list + return kept # pyright: ignore[reportReturnType] # Codex passthrough items sit outside the OpenAI input union + def map_openai_params( self, response_api_optional_params: ResponsesAPIOptionalRequestParams, diff --git a/litellm/llms/cerebras/chat.py b/litellm/llms/cerebras/chat.py index 8827b0afd87..c3aa26ade35 100644 --- a/litellm/llms/cerebras/chat.py +++ b/litellm/llms/cerebras/chat.py @@ -68,6 +68,8 @@ class CerebrasConfig(OpenAIGPTConfig): "tool_choice", "tools", "user", + "max_retries", + "extra_headers", ] # Only add reasoning_effort for models that support it diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index b7ddb55ae89..bb582f677be 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1203,6 +1203,7 @@ class BaseLLMHTTPHandler: headers=headers, data=json.dumps(request_data), timeout=timeout, + logging_obj=logging_obj, ) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index ffa3de0d5c6..3a65e4a9426 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -16,7 +16,7 @@ def _normalize_reasoning_effort_for_chat_completion( ) -> str | None: """Convert reasoning_effort to the string format expected by OpenAI chat completion API. - The chat completion API expects a simple string: 'none', 'low', 'medium', 'high', or 'xhigh'. + The chat completion API expects an effort string such as 'low' or 'high'. Config/deployments may pass the Responses API format: {'effort': 'high', 'summary': 'detailed'}. """ if value is None: diff --git a/litellm/llms/together_ai/chat.py b/litellm/llms/together_ai/chat.py deleted file mode 100644 index 58d47e45faa..00000000000 --- a/litellm/llms/together_ai/chat.py +++ /dev/null @@ -1,58 +0,0 @@ -""" -Support for OpenAI's `/v1/chat/completions` endpoint. - -Calls done in OpenAI/openai.py as TogetherAI is openai-compatible. - -Docs: https://docs.together.ai/reference/completions-1 -""" - -from typing import Final - -from litellm._logging import verbose_logger -from litellm.utils import supports_function_calling - -from ..openai.chat.gpt_transformation import OpenAIGPTConfig - - -class TogetherAIConfig(OpenAIGPTConfig): - def get_supported_openai_params(self, model: str) -> list: - """ - Only some together models support response_format / tool calling - - Docs: https://docs.together.ai/docs/json-mode - """ - # Use supports_function_calling() — which reads _get_model_info_helper - # directly — instead of get_model_info(). get_model_info() calls - # get_supported_openai_params() as its first step, which routes back - # into this method for together_ai models, creating a recursion that - # only terminates when Python's recursion limit or the "not mapped" - # exception in _get_model_info_helper is hit (~332 deep calls). - supports_fc: bool | None = None - try: - supports_fc = supports_function_calling(model, custom_llm_provider="together_ai") - except Exception as e: - verbose_logger.debug("Error getting supported openai params: %s", e) - - optional_params: Final = super().get_supported_openai_params(model) - if supports_fc is not True: - verbose_logger.debug( - "Only some together models support function calling/response_format. Docs - https://docs.together.ai/docs/function-calling" - ) - optional_params.remove("tools") - optional_params.remove("tool_choice") - optional_params.remove("function_call") - optional_params.remove("response_format") - return optional_params - - def map_openai_params( - self, - non_default_params: dict, - optional_params: dict, - model: str, - drop_params: bool, - ) -> dict: - mapped_openai_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params) - - if "response_format" in mapped_openai_params and mapped_openai_params["response_format"] == {"type": "text"}: - mapped_openai_params.pop("response_format") - return mapped_openai_params diff --git a/litellm/llms/together_ai/chat/__init__.py b/litellm/llms/together_ai/chat/__init__.py new file mode 100644 index 00000000000..f260d9126d7 --- /dev/null +++ b/litellm/llms/together_ai/chat/__init__.py @@ -0,0 +1,3 @@ +from .transformation import TogetherAIChatConfig as TogetherAIChatConfig + +TogetherAIConfig = TogetherAIChatConfig diff --git a/litellm/llms/together_ai/chat/transformation.py b/litellm/llms/together_ai/chat/transformation.py new file mode 100644 index 00000000000..3162a34f1b9 --- /dev/null +++ b/litellm/llms/together_ai/chat/transformation.py @@ -0,0 +1,89 @@ +""" +Translates from OpenAI's `/v1/chat/completions` to Together AI's `/v1/chat/completions`. + +Docs: https://docs.together.ai/docs/chat-overview +""" + +from collections.abc import Container +from types import MappingProxyType +from typing import Final + +import litellm +from litellm._logging import verbose_logger +from litellm.exceptions import UnsupportedParamsError +from litellm.utils import supports_function_calling + +from ...openai.chat.gpt_transformation import OpenAIGPTConfig + +TOOL_CALLING_PARAMS: Final = ("tools", "tool_choice", "function_call") +PLAIN_TEXT_RESPONSE_FORMAT: Final = MappingProxyType({"type": "text"}) +FUNCTION_CALLING_DOCS_URL: Final = "https://docs.together.ai/docs/function-calling" + + +def _function_calling_verdict(model: str) -> bool | None: + try: + if supports_function_calling(model, custom_llm_provider="together_ai"): + return True + except Exception as e: + verbose_logger.debug("Error checking together_ai function calling support for %s: %s", model, e) + registry_entry: Final = litellm.model_cost.get(f"together_ai/{model}") + if isinstance(registry_entry, dict) and registry_entry.get("supports_function_calling") is False: + return False + return None + + +def _tool_params_to_drop(passed_params: Container[str], model: str, drop_params: bool) -> tuple[str, ...]: + passed_tool_params: Final = tuple(param for param in TOOL_CALLING_PARAMS if param in passed_params) + if not passed_tool_params: + return () + verdict: Final = _function_calling_verdict(model) + if verdict is True: + return () + if verdict is None: + verbose_logger.warning( + "together_ai model %s has no function calling entry in the model registry; passing %s through for Together to validate. Docs - %s", + model, + ", ".join(passed_tool_params), + FUNCTION_CALLING_DOCS_URL, + ) + return () + if drop_params or litellm.drop_params: + verbose_logger.warning( + "together_ai model %s does not support function calling per the model registry; dropping %s. Docs - %s", + model, + ", ".join(passed_tool_params), + FUNCTION_CALLING_DOCS_URL, + ) + return passed_tool_params + raise UnsupportedParamsError( + status_code=500, + message=f"together_ai does not support parameters: {', '.join(passed_tool_params)}, for model={model}. To drop it from the call, set `litellm.drop_params = True`.", + ) + + +class TogetherAIChatConfig(OpenAIGPTConfig): + def get_supported_openai_params(self, model: str) -> list: + supports_fc: Final = _function_calling_verdict(model) + supported_params: Final = super().get_supported_openai_params(model) + if supports_fc is True: + return supported_params + verbose_logger.debug( + "Only some together models support response_format. Docs - https://docs.together.ai/docs/function-calling" + ) + return [ # mutable-ok: the inherited contract returns a plain list; building fresh avoids mutating the base class's value + param for param in supported_params if param != "response_format" + ] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + mapped_openai_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params) + for param in _tool_params_to_drop(mapped_openai_params, model, drop_params): + mapped_openai_params.pop(param) + if mapped_openai_params.get("response_format") == PLAIN_TEXT_RESPONSE_FORMAT: + mapped_openai_params.pop("response_format") + return mapped_openai_params diff --git a/litellm/llms/together_ai/rerank/handler.py b/litellm/llms/together_ai/rerank/handler.py index 10246451a9d..b8079e52c97 100644 --- a/litellm/llms/together_ai/rerank/handler.py +++ b/litellm/llms/together_ai/rerank/handler.py @@ -16,11 +16,16 @@ from litellm.llms.together_ai.rerank.transformation import TogetherAIRerankConfi from litellm.types.rerank import RerankRequest, RerankResponse +def _rerank_url(api_base: str) -> str: + return f"{api_base.rstrip('/')}/rerank" + + class TogetherAIRerank(BaseLLM): def rerank( self, model: str, api_key: str, + api_base: str, query: str, documents: list[str | dict[str, Any]], top_n: int | None = None, @@ -46,10 +51,10 @@ class TogetherAIRerank(BaseLLM): raise ValueError("TogetherAI does not support max_chunks_per_doc") if _is_async: - return self.async_rerank(request_data_dict, api_key) # Call async method + return self.async_rerank(request_data_dict, api_key, api_base) response: Final = client.post( - "https://api.together.xyz/v1/rerank", + _rerank_url(api_base), headers={ "accept": "application/json", "content-type": "application/json", @@ -69,11 +74,12 @@ class TogetherAIRerank(BaseLLM): self, request_data_dict: dict[str, Any], api_key: str, + api_base: str, ) -> RerankResponse: client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.TOGETHER_AI) # Use async client response: Final = await client.post( - "https://api.together.xyz/v1/rerank", + _rerank_url(api_base), headers={ "accept": "application/json", "content-type": "application/json", diff --git a/litellm/llms/vertex_ai/interactions/__init__.py b/litellm/llms/vertex_ai/interactions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/vertex_ai/interactions/transformation.py b/litellm/llms/vertex_ai/interactions/transformation.py new file mode 100644 index 00000000000..0764a8bea62 --- /dev/null +++ b/litellm/llms/vertex_ai/interactions/transformation.py @@ -0,0 +1,149 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Final + +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig +from litellm.llms.vertex_ai.common_utils import validate_vertex_location +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + +VERTEX_INTERACTIONS_API_VERSION: Final = "v1beta1" +VERTEX_INTERACTIONS_DEFAULT_LOCATION: Final = "global" + + +@dataclass(frozen=True, slots=True) +class VertexInteractionsTarget: + base_url: str + project_id: str + location: str + + @property + def collection_url(self) -> str: + return ( + f"{self.base_url}/{VERTEX_INTERACTIONS_API_VERSION}" + f"/projects/{self.project_id}/locations/{self.location}/interactions" + ) + + def interaction_url(self, interaction_id: str) -> str: + encoded_interaction_id: Final = encode_url_path_segment(interaction_id, field_name="interaction_id") + return f"{self.collection_url}/{encoded_interaction_id}" + + +class VertexAIInteractionsConfig(VertexBase, GoogleAIStudioInteractionsConfig): + def __init__( + self, + mint_access_token: Callable[[VERTEX_CREDENTIALS_TYPES | None, str | None], tuple[str, str]] | None = None, + ) -> None: + super().__init__() + self._mint_access_token: Final[Callable[[VERTEX_CREDENTIALS_TYPES | None, str | None], tuple[str, str]]] = ( + mint_access_token or self._mint_access_token_with_vertex_base + ) + + def _mint_access_token_with_vertex_base( + self, + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + ) -> tuple[str, str]: + return self._ensure_access_token( + credentials=credentials, project_id=project_id, custom_llm_provider="vertex_ai" + ) + + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.VERTEX_AI + + @property + def api_version(self) -> str: + return VERTEX_INTERACTIONS_API_VERSION + + def get_default_vertex_location(self) -> str: + return VERTEX_INTERACTIONS_DEFAULT_LOCATION + + def _mint(self, litellm_params: GenericLiteLLMParams) -> tuple[str, str]: + raw_params: Final = litellm_params.model_dump() + return self._mint_access_token( + self.safe_get_vertex_ai_credentials(raw_params), + self.safe_get_vertex_ai_project(raw_params), + ) + + def _target(self, api_base: str | None, litellm_params: GenericLiteLLMParams) -> VertexInteractionsTarget: + _, project_id = self._mint(litellm_params) + if not project_id: + raise ValueError( + "Vertex AI project is required. Set vertex_project, litellm.vertex_project, or VERTEXAI_PROJECT" + ) + location: Final = validate_vertex_location( + self.explicit_vertex_ai_location(litellm_params.model_dump()) or VERTEX_INTERACTIONS_DEFAULT_LOCATION + ) + return VertexInteractionsTarget( + base_url=self.get_api_base(api_base or None, location), + project_id=project_id, + location=location, + ) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + litellm_params: GenericLiteLLMParams | None, + ) -> dict: # mutable-ok: BaseInteractionsAPIConfig declares plain-dict headers + access_token, _ = self._mint(litellm_params or GenericLiteLLMParams()) + return { # mutable-ok: BaseInteractionsAPIConfig declares plain-dict headers + "Content-Type": "application/json", + "Authorization": f"Bearer {access_token}", + **headers, + } + + def get_complete_url( + self, + api_base: str | None, + model: str | None, + agent: str | None = None, + litellm_params: Mapping[str, object] | None = None, + stream: bool | None = None, + ) -> str: + params: Final = ( + GenericLiteLLMParams.model_validate(litellm_params) if litellm_params else GenericLiteLLMParams() + ) + collection_url: Final = self._target(api_base, params).collection_url + return f"{collection_url}?alt=sse" if stream else collection_url + + def _interaction_by_id_request( + self, + interaction_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + url_suffix: str = "", + ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body + target: Final = self._target(api_base or None, litellm_params) + return f"{target.interaction_url(interaction_id)}{url_suffix}", {} # mutable-ok: same base contract + + def transform_get_interaction_request( + self, + interaction_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body + return self._interaction_by_id_request(interaction_id, api_base, litellm_params) + + def transform_delete_interaction_request( + self, + interaction_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body + return self._interaction_by_id_request(interaction_id, api_base, litellm_params) + + def transform_cancel_interaction_request( + self, + interaction_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body + return self._interaction_by_id_request(interaction_id, api_base, litellm_params, url_suffix=":cancel") diff --git a/litellm/main.py b/litellm/main.py index d3967473f99..98f92e50599 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -24,6 +24,7 @@ from concurrent import futures from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait 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 litellm._logging import _redact_string @@ -416,7 +417,7 @@ async def acompletion( logprobs: bool | None = None, top_logprobs: int | None = None, deployment_id=None, - reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None = None, + reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "max", "default"] | None = None, verbosity: Literal["low", "medium", "high"] | None = None, safety_identifier: str | None = None, service_tier: str | None = None, @@ -602,7 +603,7 @@ async def acompletion( _, custom_llm_provider, _, _ = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider, - api_base=base_url, + api_base=kwargs.get("api_base") or base_url, ) fallbacks = fallbacks or litellm.model_fallbacks @@ -1811,6 +1812,56 @@ def _complete_fireworks_ai( return response +def _complete_together_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion: Final = ctx.acompletion + api_base: Final = ctx.api_base + api_key: Final = ctx.api_key + client: Final = _dispatch_client_http(ctx) + custom_llm_provider: Final = ctx.custom_llm_provider + headers: Final = ctx.headers + litellm_params: Final = ctx.litellm_params + logging: Final = ctx.logging + messages: Final = ctx.messages + model: Final = ctx.model + model_response: Final = ctx.model_response + optional_params: Final = ctx.optional_params + provider_config: Final = ctx.provider_config + shared_session: Final = ctx.shared_session + stream: Final = ctx.stream + timeout: Final = ctx.timeout + + try: + response: Final = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args=MappingProxyType({"headers": headers}), + ) + raise + + return response + + def _complete_heroku(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base @@ -4920,7 +4971,7 @@ def completion( logit_bias: dict | None = None, user: str | None = None, # openai v1.0+ new params - reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None = None, + reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "max", "default"] | None = None, verbosity: Literal["low", "medium", "high"] | None = None, response_format: dict | type[BaseModel] | None = None, seed: int | None = None, @@ -5600,6 +5651,8 @@ def completion( elif custom_llm_provider == "fireworks_ai": ## COMPLETION CALL response = _complete_fireworks_ai(_dispatch_ctx) + elif custom_llm_provider == "together_ai": + response = _complete_together_ai(_dispatch_ctx) elif custom_llm_provider == "heroku": response = _complete_heroku(_dispatch_ctx) @@ -5649,7 +5702,6 @@ def completion( or custom_llm_provider == "volcengine" or custom_llm_provider == "anyscale" or custom_llm_provider == "openai" - or custom_llm_provider == "together_ai" or custom_llm_provider == "nebius" or custom_llm_provider == "wandb" or custom_llm_provider == "clarifai" @@ -5699,14 +5751,6 @@ def completion( response = _complete_openrouter(_dispatch_ctx) elif custom_llm_provider == "vercel_ai_gateway": response = _complete_vercel_ai_gateway(_dispatch_ctx) - elif ( - custom_llm_provider == "together_ai" - or ("togethercomputer" in model) - or (model in litellm.together_ai_models) - ): - """ - Deprecated. We now do together ai calls via the openai client - https://docs.together.ai/docs/openai-api-compatibility - """ elif custom_llm_provider == "palm": raise ValueError( "Palm was decommisioned on October 2024. Please use the `gemini/` route for Gemini Google AI Studio Models. Announcement: https://ai.google.dev/palm_docs/palm?hl=en" diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9f953e11df1..9ca8d9e1bac 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -37886,6 +37886,7 @@ "output_cost_per_token": 1e-07 }, "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo": { + "deprecation_date": "2026-02-06", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -37902,6 +37903,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-235B-A22B-Instruct-2507-tput": { + "deprecation_date": "2026-07-10", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "max_input_tokens": 262000, @@ -37914,6 +37916,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-235B-A22B-Thinking-2507": { + "deprecation_date": "2026-04-16", "input_cost_per_token": 6.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 256000, @@ -37926,6 +37929,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-235B-A22B-fp8-tput": { + "deprecation_date": "2026-02-06", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "max_input_tokens": 40000, @@ -37937,6 +37941,7 @@ "supports_tool_choice": false }, "together_ai/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { + "deprecation_date": "2026-06-04", "input_cost_per_token": 2e-06, "litellm_provider": "together_ai", "max_input_tokens": 256000, @@ -37949,11 +37954,15 @@ "supports_tool_choice": true }, "together_ai/deepseek-ai/DeepSeek-R1": { + "deprecation_date": "2026-05-14", "input_cost_per_token": 3e-06, "litellm_provider": "together_ai", "max_input_tokens": 128000, "max_output_tokens": 20480, "max_tokens": 20480, + "metadata": { + "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro" + }, "mode": "chat", "output_cost_per_token": 7e-06, "supports_function_calling": true, @@ -37962,6 +37971,7 @@ "supports_tool_choice": true }, "together_ai/deepseek-ai/DeepSeek-R1-0528-tput": { + "deprecation_date": "2026-02-03", "input_cost_per_token": 5.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 128000, @@ -37979,6 +37989,9 @@ "max_input_tokens": 65536, "max_output_tokens": 8192, "max_tokens": 8192, + "metadata": { + "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro" + }, "mode": "chat", "output_cost_per_token": 1.25e-06, "supports_function_calling": true, @@ -37987,9 +38000,13 @@ "supports_tool_choice": true }, "together_ai/deepseek-ai/DeepSeek-V3.1": { + "deprecation_date": "2026-05-14", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_tokens": 16384, + "metadata": { + "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro" + }, "mode": "chat", "output_cost_per_token": 1.7e-06, "source": "https://www.together.ai/models/deepseek-v3-1", @@ -38001,6 +38018,7 @@ "max_output_tokens": 16384 }, "together_ai/meta-llama/Llama-3.2-3B-Instruct-Turbo": { + "deprecation_date": "2026-03-06", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -38009,16 +38027,21 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo": { - "input_cost_per_token": 8.8e-07, + "input_cost_per_token": 1.04e-06, "litellm_provider": "together_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 8.8e-07, + "output_cost_per_token": 1.04e-06, + "source": "https://docs.together.ai/docs/serverless-models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true }, "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo-Free": { + "deprecation_date": "2025-11-13", "input_cost_per_token": 0, "litellm_provider": "together_ai", "mode": "chat", @@ -38029,6 +38052,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { + "deprecation_date": "2026-03-31", "input_cost_per_token": 2.7e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -38039,6 +38063,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Llama-4-Scout-17B-16E-Instruct": { + "deprecation_date": "2026-02-06", "input_cost_per_token": 1.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -38049,6 +38074,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": { + "deprecation_date": "2026-02-06", "input_cost_per_token": 3.5e-06, "litellm_provider": "together_ai", "mode": "chat", @@ -38059,6 +38085,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": { + "deprecation_date": "2026-02-25", "input_cost_per_token": 8.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -38069,6 +38096,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": { + "deprecation_date": "2026-03-06", "input_cost_per_token": 1.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -38079,6 +38107,7 @@ "supports_tool_choice": true }, "together_ai/mistralai/Mistral-7B-Instruct-v0.1": { + "deprecation_date": "2025-11-13", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -38087,6 +38116,7 @@ "supports_tool_choice": true }, "together_ai/mistralai/Mistral-Small-24B-Instruct-2501": { + "deprecation_date": "2026-04-02", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -38094,6 +38124,7 @@ "supports_tool_choice": true }, "together_ai/mistralai/Mixtral-8x7B-Instruct-v0.1": { + "deprecation_date": "2026-04-16", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -38106,6 +38137,9 @@ "together_ai/moonshotai/Kimi-K2-Instruct": { "input_cost_per_token": 1e-06, "litellm_provider": "together_ai", + "metadata": { + "successor": "together_ai/moonshotai/Kimi-K3" + }, "mode": "chat", "output_cost_per_token": 3e-06, "source": "https://www.together.ai/models/kimi-k2-instruct", @@ -38149,6 +38183,7 @@ "supports_tool_choice": true }, "together_ai/zai-org/GLM-4.5-Air-FP8": { + "deprecation_date": "2026-04-02", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "max_input_tokens": 128000, @@ -38166,6 +38201,9 @@ "max_input_tokens": 200000, "max_output_tokens": 200000, "max_tokens": 200000, + "metadata": { + "successor": "together_ai/zai-org/GLM-5.2" + }, "mode": "chat", "output_cost_per_token": 2.2e-06, "source": "https://www.together.ai/models/glm-4-6", @@ -38175,11 +38213,15 @@ "supports_tool_choice": true }, "together_ai/zai-org/GLM-4.7": { + "deprecation_date": "2026-04-02", "input_cost_per_token": 4.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 200000, "max_output_tokens": 200000, "max_tokens": 200000, + "metadata": { + "successor": "together_ai/zai-org/GLM-5.2" + }, "mode": "chat", "output_cost_per_token": 2e-06, "source": "https://www.together.ai/models/glm-4-7", @@ -38189,11 +38231,15 @@ "supports_tool_choice": true }, "together_ai/moonshotai/Kimi-K2.5": { + "deprecation_date": "2026-05-21", "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, + "metadata": { + "successor": "together_ai/moonshotai/Kimi-K3" + }, "mode": "chat", "output_cost_per_token": 2.8e-06, "source": "https://www.together.ai/models/kimi-k2-5", @@ -38203,9 +38249,13 @@ "supports_reasoning": true }, "together_ai/moonshotai/Kimi-K2-Instruct-0905": { + "deprecation_date": "2026-03-06", "input_cost_per_token": 1e-06, "litellm_provider": "together_ai", "max_input_tokens": 262144, + "metadata": { + "successor": "together_ai/moonshotai/Kimi-K3" + }, "mode": "chat", "output_cost_per_token": 3e-06, "source": "https://www.together.ai/models/kimi-k2-0905", @@ -38214,9 +38264,13 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-Next-80B-A3B-Instruct": { + "deprecation_date": "2026-04-02", "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, + "metadata": { + "successor": "together_ai/Qwen/Qwen3.7-Plus" + }, "mode": "chat", "output_cost_per_token": 1.5e-06, "source": "https://www.together.ai/models/qwen3-next-80b-a3b-instruct", @@ -38226,9 +38280,13 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-Next-80B-A3B-Thinking": { + "deprecation_date": "2026-02-25", "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, + "metadata": { + "successor": "together_ai/Qwen/Qwen3.6-Plus" + }, "mode": "chat", "output_cost_per_token": 1.5e-06, "source": "https://www.together.ai/models/qwen3-next-80b-a3b-thinking", @@ -38238,6 +38296,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3.5-397B-A17B": { + "deprecation_date": "2026-06-29", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, @@ -38249,6 +38308,292 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/MiniMaxAI/MiniMax-M3": { + "input_cost_per_token": 3e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/Prism-ML/Ternary-Bonsai-27B": { + "input_cost_per_token": 0.0, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/Qwen/Qwen3.5-9B": { + "input_cost_per_token": 1.7e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/Qwen/Qwen3.6-Plus": { + "input_cost_per_token": 5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_reasoning": true + }, + "together_ai/Qwen/Qwen3.7-Max": { + "input_cost_per_token": 1.25e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/Qwen/Qwen3.7-Plus": { + "input_cost_per_token": 3.2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.28e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/Qwen/Qwen3.8-2.4T-A95B": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1010000, + "max_output_tokens": 1010000, + "max_tokens": 1010000, + "mode": "chat", + "output_cost_per_token": 6.25e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/arize-ai/qwen-2-1.5b-instruct": { + "input_cost_per_token": 1e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/deepseek-ai/DeepSeek-V4-Flash-0731": { + "input_cost_per_token": 1.4e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/deepseek-ai/DeepSeek-V4-Pro": { + "input_cost_per_token": 1.74e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 512000, + "max_output_tokens": 512000, + "max_tokens": 512000, + "mode": "chat", + "output_cost_per_token": 3.48e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/deepseek-ai/DeepSeek-V4-Pro-0813": { + "input_cost_per_token": 1.32e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 3.96e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/google/gemma-3n-E4B-it": { + "input_cost_per_token": 6e-08, + "litellm_provider": "together_ai", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.2e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/google/gemma-4-31B-it": { + "input_cost_per_token": 3.9e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 9.7e-07, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/intfloat/multilingual-e5-large-instruct": { + "input_cost_per_token": 2e-08, + "litellm_provider": "together_ai", + "max_input_tokens": 514, + "max_tokens": 514, + "mode": "embedding", + "output_cost_per_token": 2e-08, + "output_vector_size": 1024, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/meta-llama/Llama-Guard-4-12B": { + "input_cost_per_token": 2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/meta-models/Muse-Glimmer-30B": { + "input_cost_per_token": 3.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/moonshotai/Kimi-K2.7-Code": { + "input_cost_per_token": 9.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/moonshotai/Kimi-K3": { + "input_cost_per_token": 3e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/nvidia/nemotron-3-ultra-550b-a55b": { + "input_cost_per_token": 6e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 512288, + "max_output_tokens": 512288, + "max_tokens": 512288, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/pearl-ai/gemma-4-31b-it": { + "input_cost_per_token": 2.8e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 8.6e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/thinkingmachines/Inkling": { + "input_cost_per_token": 1e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 4.05e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/thinkingmachines/Inkling-Small": { + "input_cost_per_token": 5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/zai-org/GLM-5.2": { + "input_cost_per_token": 1.4e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048575, + "max_output_tokens": 1048575, + "max_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "tts-1": { "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", @@ -49016,12 +49361,13 @@ "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -49048,12 +49394,13 @@ "output_cost_per_token": 1.32e-05, "output_cost_per_token_above_272k_tokens": 1.98e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -49080,12 +49427,13 @@ "output_cost_per_token": 1.32e-06, "output_cost_per_token_above_272k_tokens": 1.98e-06, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 8a2ee2a3af8..4b30afb2f98 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -199,7 +199,7 @@ def llm_passthrough_route( api_key=api_key, ) - litellm_params_dict: Final = get_litellm_params(**kwargs) + litellm_params_dict: Final = get_litellm_params(api_key=api_key, api_base=api_base, **kwargs) if client is None: from litellm.llms.custom_httpx.http_handler import ( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0840d37ffa1..628e569e1b8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -815,6 +815,7 @@ class LiteLLMRoutes(enum.Enum): "/team/member_add", "/team/member_delete", "/team/member_update", + "/team/{team_id}/member/{user_id}/reset_spend", "/team/permissions_list", "/team/permissions_update", "/team/daily/activity", @@ -1287,6 +1288,16 @@ class RegenerateKeyRequest(GenerateKeyRequest): class ResetSpendRequest(LiteLLMPydanticObjectBase): reset_to: float + @field_validator("reset_to", mode="before") + @classmethod + def reject_bool_reset_to(cls, v): + # bool is a subclass of int, so pydantic silently coerces True/False into + # 1.0/0.0 for a `float` field: a caller who accidentally sends a boolean + # would otherwise get an unintended spend reset instead of a 422. + if isinstance(v, bool): + raise ValueError("reset_to must be a number, not a boolean") # noqa: TRY004 # pydantic needs ValueError + return v + class KeyRequest(LiteLLMPydanticObjectBase): keys: list[str] | None = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e7b98b3cc7f..52857909621 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -71,7 +71,6 @@ from litellm.proxy.auth.budget_throttle import ( ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation -from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, _safe_get_request_query_params, @@ -87,6 +86,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( object_permission_cache_key, tag_cache_key, tag_registry_cache_key, + team_membership_auth_cache_key, + team_membership_reservation_cache_key, ) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( @@ -1129,7 +1130,8 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: Parameters: - user_route: str - the route the user is trying to call - - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. + - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. Entries are a route group name + (e.g. "openai_routes"), an exact route, or a trailing-wildcard prefix (e.g. "/internal-models/*"). """ from starlette.routing import compile_path @@ -1139,7 +1141,7 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: regex, _, _ = compile_path(template) if regex.match(user_route): return True - elif allowed_route == user_route: + elif RouteChecks.route_matches_wildcard_pattern(route=user_route, pattern=allowed_route): return True return False @@ -1967,7 +1969,7 @@ async def get_team_membership( if user_id is None or team_id is None: return None - _key: Final = f"team_membership:{user_id}:{team_id}" + _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) # check if in cache cached_membership_obj: Final = await user_api_key_cache.async_get_cache( @@ -2402,6 +2404,116 @@ async def _cache_team_object( ) +async def invalidate_team_member_spend_state( + user_id: str, + team_id: str, + user_api_key_cache: UserApiKeyCache, + new_spend: float | None = None, +) -> None: + """ + Clear every cached read path for one team member's budget so a spend + reset or a raised cap takes effect on the next request instead of + waiting on the membership cache's TTL. + + Two independently-keyed cache entries hold the same LiteLLM_TeamMembership + row: user_api_key_auth.py's admission check writes ``{team_id}_{user_id}``, + while budget_reservation.py's pre-call reservation and auth_checks.py's own + get_team_membership() (used by _check_team_member_budget) both write + ``team_membership:{user_id}:{team_id}``. Both formats must be invalidated + explicitly; writing one does not refresh the other. All keys are also + broadcast (LIT-3803): each worker's own in-memory copy (membership object, + spend counter, or the counter's own short-TTL DB-floor marker) survives + eviction elsewhere until its TTL, so the handling worker alone clearing its + copy leaves every other worker still enforcing the pre-reset budget. + + ``new_spend`` is only passed by reset_team_member_spend_fn, which knows the + exact post-reset value: it is SET everywhere (matching /key/{key}/reset_spend's + own precedent) rather than deleted, so a worker's next read reflects it + directly instead of re-deriving it through a DB reseed. team_member_update + only changes the budget cap, not the tracked spend, so it passes no + new_spend; the live spend counter is untouched in that case (deleting it + would force a reseed from the DB's own spend column, which lags the live + counter via periodic batch writes, briefly under-enforcing the raised cap + against a spend value lower than what was actually tracked) and only the + membership caches carrying the new cap are invalidated. + + The floor marker (``spend_db_floor:``, proxy_server.py's + _authoritative_floor_spend) caches the pre-reset DB spend for + SPEND_DB_FLOOR_CACHE_TTL_SECONDS; left stale after a real reset, a request + landing on the pod that cached it can read that higher floor and raise the + counter right back above the just-reset spend. It is overwritten here with + the post-reset floor (not merely deleted) and _authoritative_floor_spend + re-checks the marker after its DB read, so a floor read already in flight + on this pod when the reset commits cannot clobber it with the pre-reset + value. Both keys are broadcast as SETs carrying new_spend, not deletes: + every subscriber (remote pods AND this pod's own, which receives its own + message) writes the post-reset value, so the self-delivered message cannot + erase the guard just written here. + + Raises HTTPException(503) if Redis still holds the stale pre-reset counter + after both the SET and the fallback DELETE fail: budget checks read Redis + first, so returning success would leave the old value authoritative for + every worker despite the DB write having committed. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, + publish_auth_cache_invalidation, + ) + + if new_spend is not None: + from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache + + spend_counter_key: Final = f"spend:team_member:{user_id}:{team_id}" + spend_db_floor_key: Final = f"spend_db_floor:{spend_counter_key}" + + spend_counter_cache.in_memory_cache.set_cache(key=spend_counter_key, value=new_spend, ttl=60) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache(key=spend_counter_key, value=new_spend, ttl=60) + except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up + verbose_proxy_logger.warning( + "Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next " + "read reseeds from the DB rather than keeping the stale pre-reset value authoritative", + spend_counter_key, + e, + ) + try: + await spend_counter_cache.redis_cache.async_delete_cache(key=spend_counter_key) + except Exception: # noqa: BLE001 # stale value now authoritative in Redis; surface instead of reporting success + verbose_proxy_logger.warning( + "Failed to delete stale spend counter %s in Redis after a failed reset write", + spend_counter_key, + exc_info=True, + ) + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail={ # mutable-ok: HTTPException.detail takes a dict + "error": "Spend was reset in the database, but Redis is unreachable and still " + "holds the pre-reset counter. Retry once Redis is reachable." + }, + ) from e + + spend_counter_cache.in_memory_cache.set_cache( + key=spend_db_floor_key, + value=new_spend, + ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS, + ) + await publish_auth_cache_invalidation(cache_key=spend_counter_key, new_value=new_spend, ttl=60) + await publish_auth_cache_invalidation( + cache_key=spend_db_floor_key, + new_value=new_spend, + ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS, + ) + + await evict_and_broadcast( + cache_keys=( + team_membership_auth_cache_key(team_id=team_id, user_id=user_id), + team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), + ), + user_api_key_cache=user_api_key_cache, + ) + + async def delete_cache_team_object( team_id: str, team_alias: str | None, @@ -2629,20 +2741,9 @@ async def _get_team_object_from_user_api_key_cache( async def _get_team_object_from_cache( key: str, - proxy_logging_obj: ProxyLogging | None, user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None, ) -> LiteLLM_TeamTableCachedObj | None: - ## INTERNAL USAGE CACHE (plain DualCache) — checked before UserApiKeyCache stores ## - if proxy_logging_obj is not None and proxy_logging_obj.internal_usage_cache.dual_cache: - cached_raw: Final = await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache( - key=key, parent_otel_span=parent_otel_span - ) - if cached_raw is not None: - from_internal: Final = CacheCodec.deserialize(cached_raw, LiteLLM_TeamTableCachedObj) - if from_internal is not None: - return from_internal - decoded: Final = await user_api_key_cache.async_get_cache( key=key, parent_otel_span=parent_otel_span, @@ -2678,7 +2779,6 @@ async def get_team_object( if not check_db_only: cached_team_obj: Final = await _get_team_object_from_cache( key=key, - proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) @@ -2841,7 +2941,6 @@ async def get_team_object_by_alias( cached_team_obj: Final = await _get_team_object_from_cache( key=cache_key, - proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 233679126f8..a42187b3a44 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -11,6 +11,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import EMPTY_MAPPING from litellm.integrations.otel.runtime import seed_request_identity +from litellm.litellm_core_utils.core_helpers import is_expected_client_error from litellm.proxy._types import ( LitellmUserRoles, ProxyErrorTypes, @@ -109,7 +110,12 @@ class UserAPIKeyAuthExceptionHandler: request=request, use_x_forwarded_for=general_settings.get("use_x_forwarded_for") is True, ) - verbose_proxy_logger.exception( + log_fn: Final = ( + verbose_proxy_logger.error + if is_expected_client_error(e) and not litellm.log_client_error_tracebacks + else verbose_proxy_logger.exception + ) + log_fn( "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s", e, requester_ip, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index d04a71535ef..9b1a6ba5aa7 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -608,7 +608,7 @@ def route_in_additonal_public_routes(current_route: str): # Check wildcard patterns for route_pattern in routes_defined: - if RouteChecks._route_matches_wildcard_pattern(route=current_route, pattern=route_pattern): + if RouteChecks.route_matches_wildcard_pattern(route=current_route, pattern=route_pattern): return True return False diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index cea21ca088b..4dba2497bb9 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -181,7 +181,7 @@ class RouteChecks: # check if wildcard pattern is allowed for allowed_route in valid_token.allowed_routes: - if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): return True if denied_auth_enforced_pass_through_route: @@ -329,7 +329,7 @@ class RouteChecks: route_allowed = True break - if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): route_allowed = True break @@ -397,7 +397,7 @@ class RouteChecks: return True # Check for wildcard patterns like "/containers/*" if RouteChecks._is_wildcard_pattern(pattern=openai_route): - if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=openai_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=openai_route): return True # Check for Google routes with placeholders like "/v1beta/models/{model_name}:generateContent" @@ -517,7 +517,7 @@ class RouteChecks: return pattern.endswith("*") @staticmethod - def _route_matches_wildcard_pattern(route: str, pattern: str) -> bool: + def route_matches_wildcard_pattern(route: str, pattern: str) -> bool: """ Check if route matches the wildcard pattern @@ -594,7 +594,7 @@ class RouteChecks: # e.g calling /anthropic/v1/messages is allowed if allowed_routes has /anthropic/* ######################################################### if any( - RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route) + RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) for allowed_route in allowed_routes if RouteChecks._is_wildcard_pattern(pattern=allowed_route) ): diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 658d176f6a7..28d76e6799c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -87,7 +87,10 @@ from litellm.proxy.common_utils.http_parsing_utils import ( populate_request_with_path_params, ) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + team_membership_auth_cache_key, +) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.utils import ( @@ -1970,8 +1973,10 @@ async def _user_api_key_auth_builder( # Check 3. Check if user is in their team budget if not skip_budget_checks and valid_token.team_member_spend is not None: - if prisma_client is not None: - _cache_key: Final = f"{valid_token.team_id}_{valid_token.user_id}" + _user_id: Final = valid_token.user_id + _team_id: Final = valid_token.team_id + if prisma_client is not None and _user_id is not None and _team_id is not None: + _cache_key: Final = team_membership_auth_cache_key(team_id=_team_id, user_id=_user_id) team_member_info = await user_api_key_cache.async_get_cache( key=_cache_key, @@ -1979,25 +1984,21 @@ async def _user_api_key_auth_builder( ) if team_member_info is None: # read from DB - _user_id: Final = valid_token.user_id - _team_id: Final = valid_token.team_id - - if _user_id is not None and _team_id is not None: - _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first( - where={ - "user_id": _user_id, - "team_id": _team_id, - }, - include={"litellm_budget_table": True}, + _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first( + where={ + "user_id": _user_id, + "team_id": _team_id, + }, + include={"litellm_budget_table": True}, + ) + if _db_member is not None: + team_member_info = LiteLLM_TeamMembership(**_db_member.dict()) + await user_api_key_cache.async_set_cache( + key=_cache_key, + value=team_member_info, + model_type=LiteLLM_TeamMembership, + ttl=5, ) - if _db_member is not None: - team_member_info = LiteLLM_TeamMembership(**_db_member.dict()) - await user_api_key_cache.async_set_cache( - key=_cache_key, - value=team_member_info, - model_type=LiteLLM_TeamMembership, - ttl=5, - ) if team_member_info is not None and team_member_info.litellm_budget_table is not None: team_member_budget: Final = team_member_info.litellm_budget_table.max_budget @@ -2013,11 +2014,16 @@ async def _user_api_key_auth_builder( max_budget=team_member_budget, ) if team_member_spend > team_member_budget: + _entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}" raise litellm.BudgetExceededError( current_cost=team_member_spend, max_budget=team_member_budget, + message=( + f"Budget has been exceeded! TeamMember={_entity_id} " + f"Current cost: {team_member_spend}, Max budget: {team_member_budget}" + ), entity_type=Litellm_EntityType.TEAM_MEMBER.value, - entity_id=f"{valid_token.user_id}:{valid_token.team_id}", + entity_id=_entity_id, ) # Check 3. If token is expired diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index dbbf9cb673e..6fd743deb44 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -33,7 +33,7 @@ from litellm.constants import ( UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, is_expected_client_error from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, @@ -1138,24 +1138,25 @@ async def open_sse_before_first_byte( ) -def _is_azure_model_router_request(model: str) -> bool: +def _is_azure_model_router_request(model: str, hidden_params: Mapping[str, object] | None = None) -> bool: """ - Check if the requested model is an Azure Model Router. + Check if a request went down the Azure Model Router route. - Azure Model Router models follow the pattern: - - azure_ai/model_router/ - - azure_ai/model-router - - model_router/ - - model-router + ``model`` here is what the *client* sent, a model group alias with no ``model_router/`` + prefix, so matching on it alone only works when the operator happened to put "model-router" + in the alias. Where the response is in hand its stamp answers this outright, so callers + should pass ``hidden_params``. Args: model: The requested model name + hidden_params: ``_hidden_params`` from the response, when the caller has it Returns: bool: True if this is an Azure Model Router request """ - model_lower: Final = model.lower() - return "model-router" in model_lower or "model_router" in model_lower + from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + + return AzureFoundryModelInfo.is_model_router_call(model=model, hidden_params=hidden_params) def _override_openai_response_model( @@ -1223,7 +1224,7 @@ def _override_openai_response_model( return # Check if this is an Azure Model Router request - if so, preserve the actual model used - if _is_azure_model_router_request(requested_model): + if _is_azure_model_router_request(requested_model, hidden_params): verbose_proxy_logger.debug( "%s: Azure Model Router detected - preserving actual model used from response instead of overriding to router model.", log_context, @@ -1379,7 +1380,12 @@ def _log_llm_api_exception(e: Exception) -> None: "litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, upstream LLM request cancelled" ) return - verbose_proxy_logger.exception("litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - %s", e) + log_fn: Final = ( + verbose_proxy_logger.error + if is_expected_client_error(e) and not litellm.log_client_error_tracebacks + else verbose_proxy_logger.exception + ) + log_fn("litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - %s", e) async def _cancel_llm_call_on_client_disconnect( diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index acdc9728390..fb2ca6372c0 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -12,6 +12,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( ) if TYPE_CHECKING: + from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -30,15 +31,24 @@ def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: @dataclass(frozen=True, slots=True) class _CacheInvalidationMessage: cache_key: str + new_value: float | None = None + ttl: float | None = None -def _cache_invalidation_message_json(cache_key: str) -> str: - return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key))) +def _cache_invalidation_message_json(cache_key: str, new_value: float | None = None, ttl: float | None = None) -> str: + message: Final = asdict(_CacheInvalidationMessage(cache_key=cache_key, new_value=new_value, ttl=ttl)) + return json.dumps({field: value for field, value in message.items() if value is not None}) -def _cache_key_from_message_data(data: object) -> str | None: +def _finite_number_or_none(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) + + +def _message_from_data(data: object) -> _CacheInvalidationMessage | None: if isinstance(data, bytes): - data = data.decode("utf-8", errors="replace") + data = data.decode("utf-8", errors="replace") # rebind-ok: normalizing the wire payload to str if not isinstance(data, str): return None try: @@ -48,14 +58,28 @@ def _cache_key_from_message_data(data: object) -> str | None: if not isinstance(parsed, dict): return None cache_key: Final = parsed.get("cache_key") - return cache_key if isinstance(cache_key, str) else None + if not isinstance(cache_key, str): + return None + return _CacheInvalidationMessage( + cache_key=cache_key, + new_value=_finite_number_or_none(parsed.get("new_value")), + ttl=_finite_number_or_none(parsed.get("ttl")), + ) -async def publish_auth_cache_invalidation(cache_key: str) -> None: +async def publish_auth_cache_invalidation( + cache_key: str, new_value: float | None = None, ttl: float | None = None +) -> None: """ Best-effort broadcast so every worker drops its local in-memory copy of a mutated management object; without this, only the handling worker and Redis are evicted and other workers keep serving the stale object until its TTL. + + Passing ``new_value`` broadcasts a SET instead of a delete: every subscriber + (including the publishing worker's own, which receives its own message) + writes the value into its additional in-memory caches rather than deleting + the key. A spend reset uses this so the handler's self-delivered message + cannot erase the freshly-written post-reset counter or floor marker. """ redis_cache: Final = coordination_redis_cache() if redis_cache is None: @@ -68,7 +92,10 @@ async def publish_auth_cache_invalidation(cache_key: str) -> None: cache_key, ) return - await client.publish(auth_cache_invalidation_channel(redis_cache), _cache_invalidation_message_json(cache_key)) + await client.publish( + auth_cache_invalidation_channel(redis_cache), + _cache_invalidation_message_json(cache_key, new_value=new_value, ttl=ttl), + ) except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e) @@ -95,15 +122,17 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us class AuthCacheInvalidationSubscriber: - __slots__ = ("_redis_cache", "_task", "_user_api_key_cache") + __slots__ = ("_additional_in_memory_caches", "_redis_cache", "_task", "_user_api_key_cache") def __init__( self, redis_cache: "RedisCache", user_api_key_cache: "UserApiKeyCache", + additional_in_memory_caches: Sequence["InMemoryCache"] = (), ) -> None: self._redis_cache = redis_cache self._user_api_key_cache = user_api_key_cache + self._additional_in_memory_caches = tuple(additional_in_memory_caches) self._task: asyncio.Task[None] | None = None def start(self) -> None: @@ -160,12 +189,18 @@ class AuthCacheInvalidationSubscriber: def _apply_message(self, message: object) -> None: data: Final = message.get("data") if isinstance(message, dict) else None - cache_key: Final = _cache_key_from_message_data(data) - if cache_key is None: + parsed: Final = _message_from_data(data) + if parsed is None: + return + if parsed.new_value is not None: + for additional_cache in self._additional_in_memory_caches: + additional_cache.set_cache(parsed.cache_key, parsed.new_value, ttl=parsed.ttl) return in_memory_cache: Final = self._user_api_key_cache.in_memory_cache if in_memory_cache is not None: - in_memory_cache.delete_cache(cache_key) + in_memory_cache.delete_cache(parsed.cache_key) + for additional_cache in self._additional_in_memory_caches: + additional_cache.delete_cache(parsed.cache_key) @staticmethod async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None: diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 93d51bdd461..b8df0105b7b 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -200,6 +200,21 @@ def end_user_restricted_registry_cache_key() -> str: return "end_user_restricted_registry" +def team_membership_auth_cache_key(team_id: str, user_id: str) -> str: + """Cache key one team member's ``LiteLLM_TeamMembership`` row is stored under for the admission check.""" + return f"{team_id}_{user_id}" + + +def team_membership_reservation_cache_key(user_id: str, team_id: str) -> str: + """Cache key the pre-call budget reservation stores the same ``LiteLLM_TeamMembership`` row under. + + Deliberately not unified with ``team_membership_auth_cache_key``: the two readers wrote independent + keys before this file existed, so a fix that invalidates one must invalidate both explicitly rather + than assume a single write is visible to both. + """ + return f"team_membership:{user_id}:{team_id}" + + def get_management_object_ttl(cache: DualCache) -> float: """ In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...). diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 2271501d480..c803d87e436 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -210,7 +210,6 @@ async def _patch_team_caches_add_access_group( for team_id in team_ids: cached_team = await _get_team_object_from_cache( key=f"team_id:{team_id}", - proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, parent_otel_span=None, ) @@ -240,7 +239,6 @@ async def _patch_team_caches_remove_access_group( for team_id in team_ids: cached_team = await _get_team_object_from_cache( key=f"team_id:{team_id}", - proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, parent_otel_span=None, ) diff --git a/litellm/proxy/management_endpoints/scim/scim_transformations.py b/litellm/proxy/management_endpoints/scim/scim_transformations.py index 2d95d0bea29..496be05b4b0 100644 --- a/litellm/proxy/management_endpoints/scim/scim_transformations.py +++ b/litellm/proxy/management_endpoints/scim/scim_transformations.py @@ -198,15 +198,8 @@ class ScimTransformations: @staticmethod def _get_scim_member_value(member: Member) -> str: - """ - Get the SCIM member value. Use user_email if available, otherwise use user_id. - SCIM member value should be the unique identifier for the user. - """ - if hasattr(member, "user_email") and member.user_email: - return member.user_email - elif hasattr(member, "user_id"): - return member.user_id or ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE - return ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE + """The member's SCIM resource id, which LiteLLM serves as user_id (RFC 7643 §8.7.1).""" + return member.user_id or ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE @staticmethod def _get_scim_member_display(member: Member) -> str: diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 7183e6cb402..6658963d024 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -176,6 +176,10 @@ class UserProvisionerHelpers: is persisted too, so re-upserting an existing email demotes a user who is no longer in the admin group instead of leaving the stale role. + IdPs like Entra manage membership exclusively through /Groups and never send + ``groups`` on POST /Users, so a request without teams means "unspecified", + not "remove from every team": existing memberships are preserved then. + Args: prisma_client: Database client new_user_request: New user request data @@ -194,7 +198,8 @@ class UserProvisionerHelpers: if not existing_user: return None - new_teams: Final = list(dict.fromkeys(new_user_request.teams or [])) + requested_teams: Final = list(dict.fromkeys(new_user_request.teams or [])) + new_teams: Final = requested_teams if requested_teams else list(existing_user.teams or []) if new_user_request.user_id != existing_user.user_id: verbose_proxy_logger.info( diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 01254d5c064..ba4fb81323f 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -16,11 +16,12 @@ import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone from types import MappingProxyType -from typing import Annotated, Final, NamedTuple, Protocol, TypedDict, TypeVar, cast +from typing import Annotated, Final, NamedTuple, NoReturn, Protocol, TypeVar, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, JsonValue +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -56,6 +57,7 @@ from litellm.proxy._types import ( PatchTeamRequest, ProxyErrorTypes, ProxyException, + ResetSpendRequest, SpecialManagementEndpointEnums, SpecialModelNames, SpecialProxyStrings, @@ -84,6 +86,7 @@ from litellm.proxy.auth.auth_checks import ( get_team_membership, get_team_object, get_user_object, + invalidate_team_member_spend_state, ) from litellm.proxy.auth.auth_utils import ( enforce_batch_enqueued_token_limit_is_admin_only, @@ -114,6 +117,7 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import ( get_daily_activity, ) from litellm.proxy.management_helpers.access_group_team_sync import ( + TEAM_ADVISORY_LOCK_SQL, AccessGroupSyncTx, invalidate_access_group_caches, reconcile_team_access_group_membership, @@ -132,6 +136,7 @@ from litellm.proxy.management_helpers.team_metadata_validation import ( validate_team_metadata_if_configured, ) from litellm.proxy.management_helpers.utils import ( + MemberWriteTx, add_new_member, management_endpoint_wrapper, ) @@ -328,11 +333,44 @@ class _TeamIdInFilter(TypedDict, total=False): team_id: Mapping[str, Sequence[str]] +class _DeletedTeamsResult(TypedDict): + deleted_teams: ReadOnly[Sequence[str]] + + +class _ErrorDetail(TypedDict): + error: ReadOnly[str] + + class _TeamCreateTx(AccessGroupSyncTx, Protocol): @property def litellm_teamtable(self) -> "_PrismaTableActions[LiteLLM_TeamTable]": ... +class _MemberDeleteTx(Protocol): + """The tables `/team/member_delete` reads while it holds the team's advisory lock. + + Reading them off the transaction keeps the whole endpoint on the one pooled connection + it already checked out: a request that has the lock but still needs another connection + can be starved by the lock waiters, which is a deadlock rather than a wait when enough + of them hold the rest of the pool.""" + + @property + def litellm_usertable(self) -> "_PrismaTableActions[LiteLLM_UserTable]": ... + + @property + def litellm_verificationtoken(self) -> "_PrismaTableActions[LiteLLM_VerificationToken]": ... + + +class _TeamDeleteTx(AccessGroupSyncTx, Protocol): + async def execute_raw(self, query: str, *args: object) -> int: ... + + @property + def litellm_teamtable(self) -> "_PrismaTableActions[LiteLLM_TeamTable]": ... + + @property + def litellm_teammembership(self) -> "_PrismaTableActions[LiteLLM_TeamMembership]": ... + + _STRIP_DELETED_TEAM_FROM_USERS_SQL: Final = """ UPDATE "LiteLLM_UserTable" SET teams = array_remove(teams, $1) WHERE $1 = ANY(teams) """ @@ -2578,8 +2616,13 @@ async def _process_team_members( prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, + tx: MemberWriteTx | None = None, ) -> tuple[list[LiteLLM_UserTable], list[LiteLLM_TeamMembership]]: - """Process and add new team members.""" + """Process and add new team members. + + ``tx`` is the caller's open transaction, when it has one, so the member writes run on the + connection it already holds instead of checking out a second one. + """ updated_users: Final[list[LiteLLM_UserTable]] = [] updated_team_memberships: Final[list[LiteLLM_TeamMembership]] = [] @@ -2605,6 +2648,7 @@ async def _process_team_members( default_team_budget_id=default_team_budget_id, allowed_models=member_allowed_models, budget_duration=data.budget_duration, + tx=tx, ) except Exception as e: raise HTTPException( @@ -2627,6 +2671,7 @@ async def _process_team_members( default_team_budget_id=default_team_budget_id, allowed_models=member_allowed_models, budget_duration=data.budget_duration, + tx=tx, ) except Exception as e: raise HTTPException( @@ -2706,65 +2751,40 @@ async def _add_team_members_to_team( user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, ) -> tuple[LiteLLM_TeamTable, list[LiteLLM_UserTable], list[LiteLLM_TeamMembership]]: - """Add team members to the team. + """Add team members to the team, under the team's advisory lock. - The members_with_roles reconciliation runs inside a transaction that locks - the team row with ``SELECT ... FOR UPDATE`` before reading the current - membership. Concurrent /team/member_add calls for the same team therefore - serialize on the row lock and each appends onto the other's committed - result, instead of both rewriting the whole JSON array from a stale - snapshot (which silently drops one member on the losing write). + The lock (``TEAM_ADVISORY_LOCK_SQL``, keyed on the team id) is taken first, and the + team is re-read under it before any write, so a delete that already committed is + visible here before this call writes anything: the user and membership writes only + happen once the re-read proves the team is still live. /team/delete takes the same + lock around its own sweep-and-delete, so the two can never interleave; whichever + acquires the lock first runs to completion before the other's re-read can proceed. - The same lock serializes this against /team/delete: the delete cannot remove - the row while the reconcile holds it, and a reconcile that finds the row - already gone cleans up after itself rather than leaving the member pointing - at a deleted team id. - """ - # Process and add new members - updated_users, updated_team_memberships = await _process_team_members( - data=data, - complete_team_data=complete_team_data, - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ) - - updated_team: Final = await _write_members_with_roles_locked( - data=data, - complete_team_data=complete_team_data, - prisma_client=prisma_client, - updated_users=updated_users, - ) - if updated_team is None: - await _sweep_deleted_team_references(team_ids=(data.team_id,), prisma_client=prisma_client) - raise HTTPException( - status_code=404, - detail={"error": f"Team={data.team_id} was deleted while this member add was running"}, - ) - - return updated_team, updated_users, updated_team_memberships - - -async def _write_members_with_roles_locked( - data: TeamMemberAddRequest, - complete_team_data: LiteLLM_TeamTable, - prisma_client: PrismaClient, - updated_users: list[LiteLLM_UserTable], -) -> LiteLLM_TeamTable | None: - """Reconcile members_with_roles under the team row lock. None when the team row is gone. - - That read is at least as recent as the user and membership writes the caller - already made, so a missing row means /team/delete committed after them. Its - post-delete sweep can have run before those writes landed, which is why the - caller sweeps this team id again rather than only reporting the 404. + The user and membership writes run on this transaction too, not on a second + connection from the pool: a lock waiter that needs a connection it hasn't got yet is + a waiter that can deadlock the pool, since enough concurrent adds for one team would + hold every connection waiting on the lock while the holder waits for a free one. """ async with prisma_client.tx() as tx: + await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, data.team_id) + locked_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, data.team_id) if locked_members is None: - return None - + gone_detail: Final[_ErrorDetail] = { + "error": f"Team={data.team_id} was deleted while this member add was running" + } + raise HTTPException(status_code=404, detail=gone_detail) complete_team_data.members_with_roles = locked_members + updated_users, updated_team_memberships = await _process_team_members( + data=data, + complete_team_data=complete_team_data, + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + tx=tx, + ) + await _update_team_members_list( data=data, complete_team_data=complete_team_data, @@ -2772,11 +2792,13 @@ async def _write_members_with_roles_locked( ) _db_team_members: Final = [m.model_dump() for m in complete_team_data.members_with_roles] - return await tx.litellm_teamtable.update( + updated_team: Final = await tx.litellm_teamtable.update( where={"team_id": data.team_id}, data={"members_with_roles": json.dumps(_db_team_members)}, ) + return updated_team, updated_users, updated_team_memberships + def _emit_team_members_metric(team: LiteLLM_TeamTable) -> None: """Update the Prometheus team members gauge after a membership change. @@ -3157,10 +3179,6 @@ async def team_member_add( litellm_proxy_admin_name=litellm_proxy_admin_name, ) - # Check if updated_team is None - if updated_team is None: - raise HTTPException(status_code=404, detail={"error": f"Team with id {data.team_id} not found"}) - _emit_team_members_metric(complete_team_data) await _create_team_member_add_audit_logs( @@ -3274,45 +3292,63 @@ async def team_member_delete( ) ## DELETE MEMBER FROM TEAM - removed_team_members, new_team_members = _cleanup_members_with_roles( - existing_team_row=existing_team_row, - data=data, - ) - - if not removed_team_members: - raise HTTPException(status_code=400, detail={"error": "User not found in team"}) - - existing_team_row.members_with_roles = new_team_members - - _db_new_team_members: Final[list[dict]] = [m.model_dump() for m in new_team_members] - - ## DELETE TEAM ID from USER ROW, IF EXISTS ## - # get user row - removed_user_ids: Final = frozenset(m.user_id for m in removed_team_members if m.user_id is not None) - key_val: Final[Mapping[str, object]] = ( - {"user_id": {"in": sorted(removed_user_ids)}} if removed_user_ids else {"user_email": data.user_email} - ) - existing_user_rows: Final[Sequence[LiteLLM_UserTable]] = await _user_db(prisma_client).find_many(where=key_val) - - # Also clean up any existing team membership rows for this user and team - user_ids_to_delete: Final = removed_user_ids.union( - (data.user_id,) if data.user_id is not None else (), - (user.user_id for user in existing_user_rows if user.user_id), - ) - - ## DELETE KEYS CREATED BY USER FOR THIS TEAM - # Fetch keys before deletion so their audit records can be persisted alongside the delete. - # An empty user_ids_to_delete still resolves cleanly: prisma's "in": [] matches no rows. - keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many( - where={ - "user_id": {"in": sorted(user_ids_to_delete)}, - "team_id": data.team_id, - } - ) - - # All four cleanups run on one connection so a failure between them leaves - # no partial removal: either every write below lands, or none of them do. + # Everything from here on runs under the team's advisory lock, the same one + # /team/member_add and /team/delete take: without it, this endpoint's own row-level + # update lock used to be the only thing serializing it against a concurrent member_add, + # and only by accident (their SELECT ... FOR UPDATE contended for the same row lock this + # UPDATE takes). Now that member_add reads under the advisory lock instead, this has to + # take it too, and re-read the roster under it rather than off the snapshot validated + # above, or a member_add that commits in between can have its addition silently + # overwritten by this delete computing from stale data. async with prisma_client.tx() as tx: + await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, data.team_id) + + fresh_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, data.team_id) + if fresh_members is None: + raise HTTPException( + status_code=400, + detail={"error": f"Team id={data.team_id} does not exist in db"}, + ) + + removed_team_members, new_team_members = _cleanup_members_with_roles( + existing_team_row=LiteLLM_TeamTable(team_id=data.team_id, members_with_roles=fresh_members), + data=data, + ) + + if not removed_team_members: + raise HTTPException(status_code=400, detail={"error": "User not found in team"}) + + existing_team_row.members_with_roles = new_team_members + + _db_new_team_members: Final[list[dict]] = [m.model_dump() for m in new_team_members] + + ## DELETE TEAM ID from USER ROW, IF EXISTS ## + # get user row + removed_user_ids: Final = frozenset(m.user_id for m in removed_team_members if m.user_id is not None) + key_val: Final[Mapping[str, object]] = ( + {"user_id": {"in": sorted(removed_user_ids)}} if removed_user_ids else {"user_email": data.user_email} + ) + member_tx: Final[_MemberDeleteTx] = tx + existing_user_rows: Final[Sequence[LiteLLM_UserTable]] = await member_tx.litellm_usertable.find_many( + where=key_val + ) + + # Also clean up any existing team membership rows for this user and team + user_ids_to_delete: Final = removed_user_ids.union( + (data.user_id,) if data.user_id is not None else (), + (user.user_id for user in existing_user_rows if user.user_id), + ) + + ## DELETE KEYS CREATED BY USER FOR THIS TEAM + # Fetch keys before deletion so their audit records can be persisted alongside the delete. + # An empty user_ids_to_delete still resolves cleanly: prisma's "in": [] matches no rows. + keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await member_tx.litellm_verificationtoken.find_many( + where={ + "user_id": {"in": sorted(user_ids_to_delete)}, + "team_id": data.team_id, + } + ) + await tx.litellm_teamtable.update( where={"team_id": data.team_id}, data={"members_with_roles": json.dumps(_db_new_team_members)}, @@ -3392,7 +3428,7 @@ async def team_member_update( Update team member budgets and team member role """ - from litellm.proxy.proxy_server import premium_user, prisma_client + from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -3491,6 +3527,12 @@ async def team_member_update( budget_patch=budget_patch, team_default_budget_id=team_default_budget_id, ) + if budget_patch: + await invalidate_team_member_spend_state( + user_id=received_user_id, + team_id=data.team_id, + user_api_key_cache=user_api_key_cache, + ) ### update team member role if data.role is not None: @@ -3527,6 +3569,125 @@ async def team_member_update( ) +def _check_not_resetting_own_spend(user_id: str, user_api_key_dict: UserAPIKeyAuth) -> None: + """ + _verify_team_access authorizes a team admin (or org admin) over their own + team, with no check that the target user_id differs from the caller. Left + unchecked, that admin could target their own LiteLLM_TeamMembership row and + repeatedly reset it to 0 right before it crosses their per-member cap, + consuming the shared team budget without the configured limit ever binding. + Only a proxy admin may reset an admin's own spend. + """ + if user_id == user_api_key_dict.user_id and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + _raise_reset_spend_error(status.HTTP_403_FORBIDDEN, "Cannot reset your own spend. Ask a proxy admin.") + + +def _raise_reset_spend_error(status_code: int, message: str) -> NoReturn: + detail: Final = {"error": message} # mutable-ok: HTTPException.detail takes a dict + raise HTTPException(status_code=status_code, detail=detail) + + +def _validate_team_member_reset_spend_value( + reset_to: object, + membership: LiteLLM_TeamMembership, +) -> float: + if not isinstance(reset_to, (int, float)): + _raise_reset_spend_error(status.HTTP_400_BAD_REQUEST, "reset_to must be a float") + + reset_to_float: Final = float(reset_to) + if not math.isfinite(reset_to_float) or reset_to_float < 0: + _raise_reset_spend_error(status.HTTP_400_BAD_REQUEST, "reset_to must be a finite number >= 0") + + current_spend: Final = membership.spend or 0.0 + if reset_to_float > current_spend: + _raise_reset_spend_error( + status.HTTP_400_BAD_REQUEST, + f"reset_to ({reset_to_float}) must be <= current spend ({current_spend})", + ) + + max_budget: Final = membership.litellm_budget_table.max_budget if membership.litellm_budget_table else None + if max_budget is not None and reset_to_float > max_budget: + _raise_reset_spend_error( + status.HTTP_400_BAD_REQUEST, + f"reset_to ({reset_to_float}) must be <= budget ({max_budget})", + ) + + return reset_to_float + + +@router.post( + "/team/{team_id}/member/{user_id}/reset_spend", + tags=["team management"], # mutable-ok: FastAPI's `tags` param is typed as list[str], not Sequence + dependencies=(Depends(user_api_key_auth),), +) +@management_endpoint_wrapper +async def reset_team_member_spend_fn( + team_id: str, + user_id: str, + data: ResetSpendRequest, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +): + """ + Reset a team member's tracked spend against their per-member budget. + + A member's spend is tracked separately from both their own personal + budget and the team's own budget (LiteLLM_TeamMembership.spend), so + neither /user/update nor /team/update can clear it: this is the only + endpoint that does. The cross-pod spend counter and cached membership + reads are invalidated so the reset takes effect on the member's next + request rather than waiting on the membership cache's TTL. + """ + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + if prisma_client is None: + _raise_reset_spend_error(status.HTTP_500_INTERNAL_SERVER_ERROR, "DB not connected. prisma_client is None") + + team_obj: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + check_db_only=True, + ) + await _verify_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict) + _check_not_resetting_own_spend(user_id=user_id, user_api_key_dict=user_api_key_dict) + + membership_where: Final = { # mutable-ok: prisma client requires a plain dict where= argument + "user_id_team_id": {"user_id": user_id, "team_id": team_id} # mutable-ok: same prisma where= argument + } + _membership_row: Final = await _team_membership_db(prisma_client).find_unique( + where=membership_where, + include={"litellm_budget_table": True}, # mutable-ok: prisma client requires a plain dict include= argument + ) + if _membership_row is None: + _raise_reset_spend_error(status.HTTP_404_NOT_FOUND, f"User {user_id} is not a member of team {team_id}.") + membership: Final = LiteLLM_TeamMembership.model_validate(_membership_row.model_dump()) + + current_spend: Final = membership.spend or 0.0 + reset_to: Final = _validate_team_member_reset_spend_value(data.reset_to, membership) + + await _team_membership_db(prisma_client).update( + where=membership_where, + data={"spend": reset_to}, # mutable-ok: prisma client requires a plain dict data= argument + ) + + await invalidate_team_member_spend_state( + user_id=user_id, + team_id=team_id, + user_api_key_cache=user_api_key_cache, + new_spend=reset_to, + ) + + return { # mutable-ok: matches this router's established untyped-response-dict convention + "team_id": team_id, + "user_id": user_id, + "spend": reset_to, + "previous_spend": current_spend, + "max_budget": membership.litellm_budget_table.max_budget if membership.litellm_budget_table else None, + } + + def _create_results_from_response( members: list[Member], response: TeamAddMemberResponse, @@ -3882,7 +4043,21 @@ async def delete_team( await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client) ## DELETE TEAMS - deleted_teams: Final = await prisma_client.delete_data(team_id_list=data.team_ids, table_name="team") + # Both the delete and the reconcile sweep run under every team's advisory lock + # (TEAM_ADVISORY_LOCK_SQL, the same one /team/member_add takes before its own writes), + # sorted so two overlapping batch deletes always request their locks in the same order. + # A member_add mid-flight for one of these teams either finishes its write and releases + # the lock before this transaction starts, in which case this sweep reaches what it wrote, + # or is still waiting on the lock, in which case its own re-read happens after this commits + # and sees the row gone before it writes anything. + delete_filter: Final[_TeamIdInFilter] = {"team_id": {"in": data.team_ids}} + async with prisma_client.tx() as tx: + for team_id in sorted(data.team_ids): + await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id) + await tx.litellm_teamtable.delete_many(where=delete_filter) + await _sweep_deleted_team_references_tx(team_ids=data.team_ids, tx=tx) + + deleted_teams: Final[_DeletedTeamsResult] = {"deleted_teams": data.team_ids} # Evict AFTER the rows are gone. Both writers of these keys (`_cache_team_object` and # `get_team_object_by_alias`) hydrate from the db, so evicting first leaves a window where a @@ -3895,12 +4070,6 @@ async def delete_team( proxy_logging_obj=proxy_logging_obj, ) - # Sweep again now the team is gone. A `/team/member_add` that landed between the first sweep - # and the delete would have re-appended the reference; an add still in flight sees the row - # missing under its own row lock and sweeps what it wrote. Both passes are idempotent, and - # keeping the first one means a failure here still leaves a team the admin can retry deleting. - await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client) - for deleted_team in team_rows: await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id) @@ -3929,6 +4098,16 @@ async def _sweep_deleted_team_references(team_ids: Sequence[str], prisma_client: _ = await _team_membership_db(prisma_client).delete_many(where=_TeamIdInFilter(team_id={"in": tuple(team_ids)})) +async def _sweep_deleted_team_references_tx(team_ids: Sequence[str], tx: _TeamDeleteTx) -> None: + """Same sweep as `_sweep_deleted_team_references`, run on the transaction that holds + every id's advisory lock and deletes the team rows, so it commits or rolls back with them.""" + for team_id in team_ids: + _ = await tx.execute_raw(_STRIP_DELETED_TEAM_FROM_USERS_SQL, team_id) + + membership_filter: Final[_TeamIdInFilter] = {"team_id": {"in": tuple(team_ids)}} + _ = await tx.litellm_teammembership.delete_many(where=membership_filter) + + async def _invalidate_deleted_key_cache( keys: Sequence[LiteLLM_VerificationToken], user_api_key_cache: UserApiKeyCache, diff --git a/litellm/proxy/management_helpers/access_group_team_sync.py b/litellm/proxy/management_helpers/access_group_team_sync.py index 55c0346e375..664e36c9f10 100644 --- a/litellm/proxy/management_helpers/access_group_team_sync.py +++ b/litellm/proxy/management_helpers/access_group_team_sync.py @@ -23,9 +23,11 @@ from pydantic import BaseModel, TypeAdapter from litellm.proxy.auth.auth_checks import _delete_cache_access_object # hashtext collisions only cost two unrelated teams a little serialization, and the -# lock is never taken by the access-group endpoints, so it cannot join their -# access-group-then-team lock order to form a cycle. -_LOCK_TEAM_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked" +# lock is never taken by the access-group endpoints as a SELECT ... FOR UPDATE row lock, +# so it cannot join their access-group-then-team lock order to form a cycle. team_endpoints +# reuses this exact statement to serialize /team/member_add and /team/delete against each +# other and against this mirror, rather than defining a second, divergent lock on the same key. +TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked" _READ_TEAM_SQL: Final = 'SELECT access_group_ids FROM "LiteLLM_TeamTable" WHERE team_id = $1' @@ -138,7 +140,7 @@ async def reconcile_team_access_group_membership(tx: AccessGroupSyncTx, team_id: concurrent write for a different team cannot be lost the way a read-modify-write of the whole array can, and the pair commits together or not at all. """ - await tx.query_raw(_LOCK_TEAM_SQL, team_id) + await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id) team_rows: Final = _TeamRows.validate_python(await tx.query_raw(_READ_TEAM_SQL, team_id)) desired: Final = (team_rows[0].access_group_ids or ()) if team_rows else () affected: Final = _AffectedGroups.validate_python(await tx.query_raw(_AFFECTED_SQL, team_id, desired)) diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index cb30ce90c7f..e2d7262fb69 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -34,7 +34,7 @@ from litellm.proxy._types import ( # key request types; user request types; tea ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import PrismaClient, jsonify_object from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.table_repositories import TeamMembershipRepository from litellm.repositories.user_repository import UserRepository @@ -79,6 +79,8 @@ class _PrismaUserTable(Protocol): self, *, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]] ) -> _PrismaUserRecord | None: ... + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_PrismaUserRecord]: ... + class _PrismaTeamMembershipTable(Protocol): """Team membership table actions the management helpers issue.""" @@ -86,6 +88,73 @@ class _PrismaTeamMembershipTable(Protocol): async def create(self, *, data: Mapping[str, object], include: Mapping[str, bool]) -> _PrismaRecord: ... +class MemberWriteTx(Protocol): + """Transaction surface `add_new_member` writes through when the caller owns one. + + A caller already holding a transaction, and with it a pooled connection plus that + transaction's locks, passes it here so these writes reuse that connection rather than + checking out another one that lock waiters may already have drained from the pool. + """ + + @property + def litellm_usertable(self) -> _PrismaUserTable: ... + + @property + def litellm_budgettable(self) -> _PrismaBudgetTable: ... + + @property + def litellm_teammembership(self) -> _PrismaTeamMembershipTable: ... + + +def _user_table(prisma_client: PrismaClient, tx: MemberWriteTx | None) -> _PrismaUserTable: + return tx.litellm_usertable if tx is not None else UserRepository(prisma_client).table + + +def _budget_table(prisma_client: PrismaClient, tx: MemberWriteTx | None) -> _PrismaBudgetTable: + return tx.litellm_budgettable if tx is not None else BudgetRepository(prisma_client).table + + +def _team_membership_table(prisma_client: PrismaClient, tx: MemberWriteTx | None) -> _PrismaTeamMembershipTable: + return tx.litellm_teammembership if tx is not None else TeamMembershipRepository(prisma_client).table + + +async def _find_users_by_email( + prisma_client: PrismaClient, tx: MemberWriteTx | None, user_email: str +) -> Sequence[_PrismaUserRecord]: + if tx is not None: + return await tx.litellm_usertable.find_many(where={"user_email": user_email}) + rows: Final[Sequence[_PrismaUserRecord] | None] = await prisma_client.get_data( + key_val={"user_email": user_email}, + table_name="user", + query_type="find_all", + ) + return rows if rows is not None else () + + +async def _upsert_user_row( + user_table: _PrismaUserTable, user_id: str, create_data: Mapping[str, object] +) -> _PrismaUserRecord | None: + """Insert the user row if it is absent, leaving an existing row as it is. + + Upserting keeps concurrent provisioning of the same new user from racing on create. + The update branch re-states user_id rather than being empty because Prisma only + compiles an upsert down to INSERT ... ON CONFLICT when the update is non-empty, and + otherwise falls back to a racy SELECT-then-INSERT. + """ + return await user_table.upsert( + where={"user_id": user_id}, + data={"create": create_data, "update": {"user_id": user_id}}, + ) + + +async def _create_user_row( + prisma_client: PrismaClient, tx: MemberWriteTx | None, user_data: dict[str, object] +) -> _PrismaUserRecord | None: + if tx is not None: + return await _upsert_user_row(tx.litellm_usertable, str(user_data["user_id"]), jsonify_object(user_data)) + return await prisma_client.insert_data(data=user_data, table_name="user") + + def get_new_internal_user_defaults(user_id: str, user_email: str | None = None) -> dict[str, object]: user_info: Final = litellm.default_internal_user_params or {} @@ -206,6 +275,7 @@ async def _clone_team_default_budget_for_member( user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, budget_duration_override: str | None = None, + tx: MemberWriteTx | None = None, ) -> str | None: """ Create a new budget row that copies the values from the team's default @@ -220,7 +290,7 @@ async def _clone_team_default_budget_for_member( member while keeping the default's other limits, so an admin can set a member's reset cadence without discarding the team default's max_budget. """ - budget_table: Final[_PrismaBudgetTable] = BudgetRepository(prisma_client).table + budget_table: Final[_PrismaBudgetTable] = _budget_table(prisma_client, tx) default_budget: Final = await budget_table.find_unique(where={"budget_id": default_team_budget_id}) if default_budget is None: return None @@ -248,7 +318,7 @@ async def _clone_team_default_budget_for_member( if cloned_data.get("budget_duration"): cloned_data["budget_reset_at"] = get_budget_reset_time(cloned_data["budget_duration"]) - new_budget: Final[_PrismaBudgetRecord] = await BudgetRepository(prisma_client).table.create(data=cloned_data) + new_budget: Final[_PrismaBudgetRecord] = await budget_table.create(data=cloned_data) return new_budget.budget_id @@ -260,6 +330,7 @@ async def _resolve_member_budget_id( allowed_models: list[str] | None, budget_duration: str | None, default_team_budget_id: str | None, + tx: MemberWriteTx | None = None, ) -> str | None: """ Resolve the budget a new team member should be linked to. @@ -279,6 +350,7 @@ async def _resolve_member_budget_id( user_api_key_dict=user_api_key_dict, litellm_proxy_admin_name=litellm_proxy_admin_name, budget_duration_override=budget_duration, + tx=tx, ) if not has_explicit_limit and budget_duration is None: @@ -295,12 +367,14 @@ async def _resolve_member_budget_id( if budget_duration is not None: budget_data["budget_duration"] = budget_duration budget_data["budget_reset_at"] = get_budget_reset_time(budget_duration=budget_duration) - budget_table: Final[_PrismaBudgetTable] = BudgetRepository(prisma_client).table + budget_table: Final[_PrismaBudgetTable] = _budget_table(prisma_client, tx) response: Final = await budget_table.create(data=budget_data) return response.budget_id -async def _append_team_id_if_absent(prisma_client: PrismaClient, user_id: str, team_id: str) -> None: +async def _append_team_id_if_absent( + prisma_client: PrismaClient, user_id: str, team_id: str, tx: MemberWriteTx | None = None +) -> None: """Append team_id to a user's teams array, only if it is not already present. The row-level filter makes the append a no-op once the team is present, so @@ -309,7 +383,7 @@ async def _append_team_id_if_absent(prisma_client: PrismaClient, user_id: str, t number of teams a user belongs to). Teams added concurrently for a different team id are unaffected, since each update filters on its own team id. """ - user_table: Final[_PrismaUserTable] = UserRepository(prisma_client).table + user_table: Final[_PrismaUserTable] = _user_table(prisma_client, tx) await user_table.update_many( where={"user_id": user_id, "NOT": {"teams": {"has": team_id}}}, data={"teams": {"push": [team_id]}}, @@ -326,6 +400,7 @@ async def add_new_member( default_team_budget_id: str | None = None, allowed_models: list[str] | None = None, budget_duration: str | None = None, + tx: MemberWriteTx | None = None, ) -> tuple[LiteLLM_UserTable, LiteLLM_TeamMembership | None]: """ Add a new member to a team @@ -334,49 +409,41 @@ async def add_new_member( - add team member w/ budget to team member table Returns created/existing user + team membership w/ budget id + + Callers already inside a transaction pass it as ``tx`` so every write here runs on that + connection instead of borrowing more from the pool while the caller's locks are held. """ returned_user: LiteLLM_UserTable | None = None returned_team_membership: LiteLLM_TeamMembership | None = None ## ADD TEAM ID, to USER TABLE IF NEW ## if new_member.user_id is not None: new_user_defaults = get_new_internal_user_defaults(user_id=new_member.user_id) - # Upsert ensures the user row exists atomically (no create race when the - # same new user is provisioned concurrently), seeding teams on create. - # The teams append lives in the filtered update below rather than the - # upsert's update branch so an already-existing user does not get a - # duplicate team id. The update branch still has to write something: - # Prisma only compiles an upsert down to INSERT ... ON CONFLICT when it - # is non-empty, and falls back to a racy SELECT-then-INSERT when it is - # not, so this re-states user_id as a no-op rather than being empty. - user_table: Final[_PrismaUserTable] = UserRepository(prisma_client).table - _returned_user: _PrismaUserRecord | None = await user_table.upsert( - where={"user_id": new_member.user_id}, - data={ - "create": {"teams": [team_id], **new_user_defaults}, - "update": {"user_id": new_member.user_id}, - }, + # The teams append lives in the filtered update below rather than the upsert's + # update branch so an already-existing user does not get a duplicate team id. + _returned_user: _PrismaUserRecord | None = await _upsert_user_row( + _user_table(prisma_client, tx), + new_member.user_id, + {"teams": [team_id], **new_user_defaults}, ) - await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id) + await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id, tx) if _returned_user is not None: returned_user = LiteLLM_UserTable.model_validate(_returned_user.model_dump()) elif new_member.user_email is not None: new_user_defaults = get_new_internal_user_defaults(user_id=str(uuid.uuid4()), user_email=new_member.user_email) ## user email is not unique acc. to prisma schema -> future improvement ### for now: check if it exists in db, if not - insert it - existing_user_row: Final[list[_PrismaUserRecord] | None] = await prisma_client.get_data( - key_val={"user_email": new_member.user_email}, - table_name="user", - query_type="find_all", + existing_user_row: Final[Sequence[_PrismaUserRecord]] = await _find_users_by_email( + prisma_client, tx, new_member.user_email ) - if existing_user_row is None or (isinstance(existing_user_row, list) and len(existing_user_row) == 0): + if len(existing_user_row) == 0: new_user_defaults["teams"] = [team_id] - _returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user") + _returned_user = await _create_user_row(prisma_client, tx, new_user_defaults) if _returned_user is not None: returned_user = LiteLLM_UserTable.model_validate(_returned_user.model_dump()) elif len(existing_user_row) == 1: user_info: Final = existing_user_row[0] - await _append_team_id_if_absent(prisma_client, user_info.user_id, team_id) + await _append_team_id_if_absent(prisma_client, user_info.user_id, team_id, tx) returned_user = LiteLLM_UserTable.model_validate(user_info.model_dump()) elif len(existing_user_row) > 1: raise HTTPException( @@ -392,10 +459,11 @@ async def add_new_member( allowed_models=allowed_models, budget_duration=budget_duration, default_team_budget_id=default_team_budget_id, + tx=tx, ) if _budget_id and returned_user is not None and returned_user.user_id is not None: - membership_table: Final[_PrismaTeamMembershipTable] = TeamMembershipRepository(prisma_client).table + membership_table: Final[_PrismaTeamMembershipTable] = _team_membership_table(prisma_client, tx) _returned_team_membership: Final = await membership_table.create( data={ "team_id": team_id, diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py index f66dc4e7bbe..001e4115374 100644 --- a/litellm/proxy/policy_engine/policy_matcher.py +++ b/litellm/proxy/policy_engine/policy_matcher.py @@ -30,7 +30,7 @@ class PolicyMatcher: """ Check if a value matches any of the given patterns. - Uses the existing RouteChecks._route_matches_wildcard_pattern helper. + Uses the existing RouteChecks.route_matches_wildcard_pattern helper. Args: value: The value to check (e.g., team alias, key alias, model) @@ -45,7 +45,7 @@ class PolicyMatcher: for pattern in patterns: # Use existing wildcard pattern matching helper - if RouteChecks._route_matches_wildcard_pattern(route=value, pattern=pattern): + if RouteChecks.route_matches_wildcard_pattern(route=value, pattern=pattern): return True return False diff --git a/litellm/proxy/policy_engine/policy_resolve_endpoints.py b/litellm/proxy/policy_engine/policy_resolve_endpoints.py index 346586c1e5a..70b98933d0f 100644 --- a/litellm/proxy/policy_engine/policy_resolve_endpoints.py +++ b/litellm/proxy/policy_engine/policy_resolve_endpoints.py @@ -100,7 +100,7 @@ def _filter_keys_by_tags(keys: list, tag_patterns: list) -> tuple: key_alias = key.key_alias or "" key_tags = _get_tags_from_metadata(key.metadata, getattr(key, "metadata_json", None)) if key_tags and any( - RouteChecks._route_matches_wildcard_pattern(route=tag, pattern=pat) + RouteChecks.route_matches_wildcard_pattern(route=tag, pattern=pat) for tag in key_tags for pat in tag_patterns ): @@ -123,7 +123,7 @@ def _filter_teams_by_tags(teams: list, tag_patterns: list) -> tuple: team_alias = team.team_alias or "" team_tags = _get_tags_from_metadata(team.metadata) if team_tags and any( - RouteChecks._route_matches_wildcard_pattern(route=tag, pattern=pat) + RouteChecks.route_matches_wildcard_pattern(route=tag, pattern=pat) for tag in team_tags for pat in tag_patterns ): @@ -152,7 +152,7 @@ async def _find_affected_by_team_patterns( for team in all_teams: team_alias = team.team_alias or "" if team_alias and any( - RouteChecks._route_matches_wildcard_pattern(route=team_alias, pattern=pat) for pat in team_patterns + RouteChecks.route_matches_wildcard_pattern(route=team_alias, pattern=pat) for pat in team_patterns ): if team_alias not in existing_teams: new_teams.append(team_alias) @@ -190,7 +190,7 @@ async def _find_affected_keys_by_alias(prisma_client: object, key_patterns: list for key in keys: key_alias = key.key_alias or "" if key_alias and any( - RouteChecks._route_matches_wildcard_pattern(route=key_alias, pattern=pat) for pat in key_patterns + RouteChecks.route_matches_wildcard_pattern(route=key_alias, pattern=pat) for pat in key_patterns ): if key_alias not in existing_keys: affected.append(key_alias) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7dced4e26b6..0abcdeaf3f6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2555,6 +2555,12 @@ async def _authoritative_floor_spend( if db_spend is None: return None + # a spend reset that committed during the DB read above wrote the post-reset + # floor to the marker; keep it over this read's now-stale pre-commit value + rechecked: Final = spend_counter_cache.in_memory_cache.get_cache(key=marker_key) + if rechecked is not None: + return float(rechecked) + spend_counter_cache.in_memory_cache.set_cache( key=marker_key, value=db_spend, @@ -6798,6 +6804,7 @@ class ProxyConfig: subscriber: Final = AuthCacheInvalidationSubscriber( redis_cache=redis_cache, user_api_key_cache=user_api_key_cache, + additional_in_memory_caches=(spend_counter_cache.in_memory_cache,), ) self.auth_cache_invalidation_subscriber = subscriber subscriber.start() diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py index 45b190c1f9d..dd5803796b7 100644 --- a/litellm/proxy/rerank_endpoints/endpoints.py +++ b/litellm/proxy/rerank_endpoints/endpoints.py @@ -90,12 +90,15 @@ async def rerank( fastapi_response.headers.update( ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, + call_id=hidden_params.get("litellm_call_id", None) or data.get("litellm_call_id", None), model_id=model_id, cache_key=cache_key, api_base=api_base, version=version, + response_cost=hidden_params.get("response_cost", None), model_region=getattr(user_api_key_dict, "allowed_model_region", ""), request_data=data, + hidden_params=hidden_params, **additional_headers, ) ) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index ce6c9330620..149f9b960a1 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -25,7 +25,11 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_utils import get_model_from_request from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key, tag_cache_key +from litellm.proxy.common_utils.user_api_key_cache import ( + end_user_cache_key, + tag_cache_key, + team_membership_reservation_cache_key, +) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -546,7 +550,9 @@ async def _get_team_member_budget_counter( if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None: return None - membership_cache_key: Final = f"team_membership:{valid_token.user_id}:{team_object.team_id}" + membership_cache_key: Final = team_membership_reservation_cache_key( + user_id=valid_token.user_id, team_id=team_object.team_id + ) cached_team_membership: Final = await user_api_key_cache.async_get_cache(key=membership_cache_key) team_membership: LiteLLM_TeamMembership | None = None if isinstance(cached_team_membership, LiteLLM_TeamMembership): diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index b6f695db512..d931d712a92 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -444,7 +444,9 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs or None ) raw_model: Final = cast(str, kwargs.get("model") or "") - model_name: Final = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) + model_name: Final = ( + standard_logging_payload.get("model") if standard_logging_payload is not None else None + ) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) try: payload: Final[SpendLogsPayload] = SpendLogsPayload( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 86d954c0913..2bbdad9c487 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -91,7 +91,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert -from litellm.litellm_core_utils.core_helpers import coerce_token_limit +from litellm.litellm_core_utils.core_helpers import coerce_token_limit, is_expected_client_error from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -2575,20 +2575,36 @@ class ProxyLogging: api_key="", ) - # log the custom exception - await litellm_logging_obj.async_failure_handler( - exception=original_exception, - traceback_exception=traceback.format_exc(), + await self._dispatch_proxy_only_failure_handlers( + litellm_logging_obj=litellm_logging_obj, + original_exception=original_exception, ) - threading.Thread( - target=litellm_logging_obj.failure_handler, - args=( - original_exception, - traceback.format_exc(), - ), - daemon=True, - ).start() + @staticmethod + async def _dispatch_proxy_only_failure_handlers( + litellm_logging_obj: Logging, + original_exception: Exception | None, + ) -> None: + """Runs the async failure handler plus the threaded sync handler. Expected + client (4xx) errors skip traceback formatting unless + litellm.log_client_error_tracebacks is set.""" + include_traceback: Final = litellm.log_client_error_tracebacks or not is_expected_client_error( + original_exception + ) + traceback_str: Final = traceback.format_exc() if include_traceback else "" + await litellm_logging_obj.async_failure_handler( + exception=original_exception, + traceback_exception=traceback_str, + ) + + threading.Thread( + target=litellm_logging_obj.failure_handler, + args=( + original_exception, + traceback_str, + ), + daemon=True, + ).start() async def post_call_success_hook( self, diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index 7efd32288e4..d636592f925 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -58,19 +58,22 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): return LiteLLM_TeamTable.model_validate(data) async def get_members_with_roles_locked(self, tx: "Prisma", team_id: str) -> list[Member] | None: - """Return the team's members_with_roles, locking the row FOR UPDATE. + """Return the team's members_with_roles. The caller must already hold + ``TEAM_ADVISORY_LOCK_SQL`` for this team_id on ``tx`` before calling this. - ``None`` when the team row is gone, which a caller holding the lock can - only see if a delete committed under it, as opposed to ``[]`` for a team - that simply has no members. + ``None`` when the team row is gone, which is only possible under that lock if + a delete committed before this read, as opposed to ``[]`` for a team that + simply has no members. - Must be called inside a transaction so the row lock is held until - commit. This serializes concurrent membership writers on the team row - so the losing writer appends onto the winner's committed result instead - of overwriting it from a stale snapshot. + A plain read is enough here because the advisory lock, not a row lock, is what + serializes this against a concurrent writer: ``SELECT ... FOR UPDATE`` would + additionally take a row lock on ``LiteLLM_TeamTable``, and the access-group + endpoints lock an access group and then a team row, so a team-row-first lock + here can deadlock with them. The advisory lock cannot, since those endpoints + never take it. """ rows: Final = await tx.query_raw( - 'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = $1 FOR UPDATE', + 'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = $1', team_id, ) if not rows: diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 15a6f18a6bb..c8f7842aebf 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -277,6 +277,8 @@ def rerank( if api_key is None: raise ValueError("TogetherAI API key is required, please set 'TOGETHERAI_API_KEY' in your environment") + api_base = dynamic_api_base or optional_params.api_base or litellm.api_base or "https://api.together.ai/v1" + response = together_rerank.rerank( model=model, query=query, @@ -286,6 +288,7 @@ def rerank( return_documents=return_documents, max_chunks_per_doc=max_chunks_per_doc, api_key=api_key, + api_base=api_base, _is_async=_is_async, ) elif _custom_llm_provider == litellm.LlmProviders.JINA_AI: diff --git a/litellm/router.py b/litellm/router.py index 6ee474730c9..d07effd0d90 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8,6 +8,7 @@ # Thank you ! We ❤️ you! - Krrish & Ishaan import asyncio +import contextlib import copy import enum import hashlib @@ -20,7 +21,7 @@ import time import traceback import weakref from collections import defaultdict -from collections.abc import AsyncGenerator, Callable, Generator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping, Sequence from functools import lru_cache from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast @@ -168,6 +169,11 @@ from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( PromptCachingDeploymentCheck, ) +from litellm.router_utils.reasoning_effort_capability import ( + deployment_is_catalog_mapped, + intersect_supported_reasoning_efforts, + resolve_supported_reasoning_efforts, +) from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_failures_for_current_minute, increment_deployment_successes_for_current_minute, @@ -243,6 +249,7 @@ from .router_utils.pattern_match_deployments import PatternMatchRouter if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from litellm.exceptions import MidStreamFallbackError from litellm.responses.streaming_iterator import ( BaseResponsesAPIStreamingIterator, ) @@ -259,6 +266,9 @@ if TYPE_CHECKING: from litellm.router_strategy.quality_router.quality_router import ( QualityRouter, ) + from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, + ) from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject from litellm.types.llms.openai import ( ResponseAPIUsage, @@ -356,6 +366,101 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) return False +# Router._aanthropic_messages_streaming_iterator buffers lifecycle chunks +# until real content commits the primary stream; a hostile or slow-starting +# upstream that never emits content or an error could otherwise grow that +# buffer without bound, so hitting this cap forces an early commit instead. +MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS: Final = 200 + + +def _anthropic_stream_should_drop_pre_content_ping(chunk: object, has_generated_content: bool) -> bool: + """A `ping` keepalive seen before any real content is dropped outright - it recurs indefinitely on a + slow-starting connection and carries nothing worth buffering toward a possible fallback.""" + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import is_anthropic_ping_chunk + + if has_generated_content: + return False + return is_anthropic_ping_chunk(chunk) + + +def _is_retriable_anthropic_status(status_code: int) -> bool: + return status_code == 429 or status_code >= 500 + + +def _anthropic_stream_should_decline_fallback(has_generated_content: bool, error: "MidStreamFallbackError") -> bool: + """ + A MidStreamFallbackError raised directly by the source iterator (the + completion-bridge path's CustomStreamWrapper, e.g. on a transport drop) + carries its own pre_first_chunk bookkeeping - gated the same way a + detected SSE error event is, so a fallback is never appended after real + content already reached the client on either path. + """ + return has_generated_content or not error.is_pre_first_chunk + + +def _anthropic_stream_commits_now(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool: + """ + Whether `chunk` should make Router._aanthropic_messages_streaming_iterator + commit to the primary Anthropic stream (real content arrived, or the + pre-content buffer cap was hit) rather than keep buffering lifecycle + frames toward a possible fallback. + """ + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + is_anthropic_content_delta_chunk, + ) + + if has_generated_content: + return False + return is_anthropic_content_delta_chunk(chunk) or buffered_chunk_count >= MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + + +class FallbackAwareAnthropicMessagesStream: + """ + Bare async generators can't carry the `_hidden_params` attribute the + proxy reads response headers off of (see + router_utils.add_retry_fallback_headers.get_hidden_params_dict), so this + thin wrapper carries it through from the source iterator - mirrors + AnthropicMessagesStreamingResponse. Used by + Router._aanthropic_messages_streaming_iterator. + """ + + def __init__(self, async_generator: AsyncGenerator[bytes, None], source_iterator: object) -> None: + self._async_generator = async_generator + self._hidden_params = dict( # mutable-ok: mutated in place by merge_fallback_hidden_params + getattr(source_iterator, "_hidden_params", None) or {} + ) + + def __aiter__(self) -> "FallbackAwareAnthropicMessagesStream": + return self + + async def __anext__(self) -> bytes: + return await self._async_generator.__anext__() + + async def aclose(self) -> None: + await self._async_generator.aclose() + + def merge_fallback_hidden_params( + self, + fallback_hidden_params: Mapping[str, object], + fallback_headers: Mapping[str, object], + ) -> None: + """ + Raw bytes can't carry their own _hidden_params the way a + ModelResponseStream/ResponsesAPI event can, so a mid-stream + fallback's provider headers (e.g. Bedrock's x-amzn-requestid) are + merged onto the wrapper itself instead - mirrors + Router._apply_fallback_hidden_params_to_item's merge shape. + """ + existing_headers: Final = cast( # cast-ok: additional_headers is always a dict[str, object] when present + "dict[str, object]", self._hidden_params.get("additional_headers") or {} + ) + self._hidden_params = { # mutable-ok: matches _hidden_params' existing dict[str, object] shape + **self._hidden_params, + **fallback_hidden_params, + "additional_headers": {**existing_headers, **fallback_headers}, # mutable-ok: same shape + } + + class RoutingArgs(enum.Enum): ttl = 60 # 1min (RPM/TPM expire key) @@ -4801,6 +4906,264 @@ class Router: ) return response + async def _aanthropic_messages_streaming_iterator( + self, + response: AsyncIterator[bytes], + initial_kwargs: dict[str, Any], # mutable-ok: mutated in-place before re-entering the fallback chain + ) -> AsyncIterator[bytes]: + """ + Wrap an anthropic_messages (/v1/messages) streaming response so a + mid-stream provider error triggers the Router's fallback chain + (parity with _acompletion_streaming_iterator for the + chat-completions path). See #24004. + + anthropic_messages goes through _ageneric_api_call_with_fallbacks + rather than _acompletion, so the returned byte iterator is never + wrapped by the chat-completions fallback handler. Two failure + shapes land here: + - the completion-bridge path (deployments with no native + /v1/messages endpoint, via + LiteLLMMessagesToCompletionTransformationHandler) already + raises MidStreamFallbackError out of its underlying + CustomStreamWrapper; this wrapper only needs to catch it. + - a native Anthropic/Bedrock passthrough never raises anything + for a provider SSE `event: error` frame (e.g. `overloaded_error`, + `internal_server_error`) - it is forwarded to the client as-is - + so this wrapper detects it via parse_anthropic_error_event and + raises MidStreamFallbackError itself. + + Only an error before any real content (a content_block_delta frame) + has reached the caller triggers a fallback attempt, mirroring the + restriction _acompletion_streaming_iterator applies: once generated + output has already reached the caller, retrying would start a + second, overlapping Anthropic message lifecycle on the same SSE + stream, so the error is left to propagate instead of being retried + invisibly. A non-retriable client error (4xx other than 429) is + never worth a fallback attempt either, so it is also left to + propagate. + + Lifecycle/bookkeeping frames (message_start, content_block_start, + ping, ...) do not by themselves disqualify a fallback attempt - + Anthropic routinely sends message_start before an overload error - + but they are BUFFERED rather than forwarded immediately, since + forwarding one and then appending a fallback attempt's own + message_start would produce two overlapping message lifecycles on + one SSE stream. Buffered frames are flushed, in order, the moment + real content arrives (the primary attempt has committed by then + anyway) or once the stream ends without ever producing content or + an error. + """ + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + aclose_if_supported, + parse_anthropic_error_event, + ) + + source_iterator: Final = response + + async def stream_with_fallbacks() -> AsyncGenerator[bytes, None]: + from litellm.exceptions import MidStreamFallbackError + + # Lifecycle/bookkeeping frames (message_start, content_block_start, + # ping, ...) are held back rather than forwarded immediately: + # Anthropic routinely sends message_start before an overload + # error, and once a byte reaches the client a fallback attempt + # can only append its OWN message_start, producing two + # overlapping message lifecycles on one SSE stream. Buffered + # frames are flushed the moment real content (content_block_delta) + # arrives - at that point the primary attempt has committed and a + # clean retry is no longer possible anyway - or once the primary + # stream ends without ever producing content. A `ping` keepalive + # is dropped outright rather than buffered, since it can recur + # indefinitely on a slow-starting connection and carries nothing + # worth preserving; hitting MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + # forces the same early commit as real content arriving, so a + # hostile or pathological upstream can't grow the buffer forever. + has_generated_content = False # rebind-ok: set once real content is seen, or the buffer cap is hit + buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline + model: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group + try: + async for chunk in source_iterator: + if _anthropic_stream_should_drop_pre_content_ping(chunk, has_generated_content): + continue + if _anthropic_stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)): + has_generated_content = True # rebind-ok: real content seen, or the buffer cap was hit + error_event = parse_anthropic_error_event(chunk) + retriable_pending_error = ( # rebind-ok: freshly computed each iteration, never carried over + not has_generated_content + and error_event is not None + and _is_retriable_anthropic_status(error_event[2]) + ) + if not has_generated_content and not retriable_pending_error and error_event is None: + buffered_lifecycle_chunks = (*buffered_lifecycle_chunks, chunk) + continue + if retriable_pending_error: + assert error_event is not None # guard-ok: retriable_pending_error implies this + _error_type, message, status_code = error_event + raise MidStreamFallbackError( + message=message, + model=model, + llm_provider="anthropic", + original_exception=litellm.exceptions.APIError( + status_code=status_code, + message=message, + llm_provider="anthropic", + model=model, + ), + is_pre_first_chunk=True, + ) + for buffered_chunk in buffered_lifecycle_chunks: + yield buffered_chunk + buffered_lifecycle_chunks = () + yield chunk + for buffered_chunk in buffered_lifecycle_chunks: + yield buffered_chunk + except MidStreamFallbackError as e: + if _anthropic_stream_should_decline_fallback(has_generated_content, e): + for buffered_chunk in buffered_lifecycle_chunks: + yield buffered_chunk + if e.original_exception is not None: + raise e.original_exception from e + raise + async for item in self._aanthropic_messages_fallback_attempt(e, initial_kwargs, wrapper): + yield item + finally: + with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): + await aclose_if_supported(source_iterator) + + # Referenced by stream_with_fallbacks via closure - assigned here, before + # the generator body ever runs, so the reference resolves fine despite + # being defined textually after the function that captures it. + wrapper: Final = FallbackAwareAnthropicMessagesStream(stream_with_fallbacks(), source_iterator) + return wrapper + + async def _aanthropic_messages_fallback_attempt( + self, + e: "MidStreamFallbackError", + initial_kwargs: dict[str, Any], # mutable-ok: mutated in-place before re-entering the fallback chain + wrapper: "FallbackAwareAnthropicMessagesStream", + ) -> AsyncGenerator[bytes, None]: + """ + Re-enters the Router's fallback chain for a mid-stream + anthropic_messages error and yields whatever the fallback attempt + produces. Split out of _aanthropic_messages_streaming_iterator to + keep each function's cyclomatic complexity within the repo's C901 + budget. + """ + from litellm.exceptions import MidStreamFallbackError + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + aclose_if_supported, + anthropic_messages_response_as_sse_events, + ) + + fallback_response = None # rebind-ok: pre-init so finally can close it if a fallback was actually attempted + try: + model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: model group + fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the common_utils list|None param + "fallbacks", self.fallbacks + ) + context_window_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below + "context_window_fallbacks", self.context_window_fallbacks + ) + content_policy_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below + "content_policy_fallbacks", self.content_policy_fallbacks + ) + initial_kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_helper + self._update_kwargs_before_fallbacks( + model=model_group, + kwargs=initial_kwargs, + metadata_variable_name="litellm_metadata", + ) + fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success + e=e, + disable_fallbacks=False, + fallbacks=fallbacks, + context_window_fallbacks=context_window_fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + model_group=model_group, + args=(), + kwargs=initial_kwargs, + include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, + ) + fallback_hidden_params, fallback_headers = Router._prepare_fallback_hidden_params(fallback_response) + wrapper.merge_fallback_hidden_params(fallback_hidden_params, fallback_headers) + if hasattr(fallback_response, "__aiter__"): + async for fallback_item in fallback_response: + yield fallback_item + else: + # A fallback can resolve to a complete AnthropicMessagesResponse + # dict even for a streaming request (e.g. an agentic tool-use + # interception loop) - yielding it as-is would put a raw dict + # into a byte stream, so it's synthesized into the SSE + # lifecycle a real stream would have sent instead. + for event in anthropic_messages_response_as_sse_events( + cast("AnthropicMessagesResponse", fallback_response) # cast-ok: non-streaming shape by elimination + ): + yield event + except Exception as fallback_error: + verbose_router_logger.error("Anthropic messages streaming fallback also failed: %s", fallback_error) + if isinstance(fallback_error, MidStreamFallbackError) and fallback_error.original_exception is not None: + raise fallback_error.original_exception from fallback_error + raise + finally: + if fallback_response is not None: + with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): + await aclose_if_supported(fallback_response) + + async def _aanthropic_messages_with_streaming_fallbacks( + self, + original_function: Callable, + **kwargs: object, # kwargs-ok: forwarded verbatim to original_function, shape varies per call site + ) -> Union["AnthropicMessagesResponse", AsyncIterator[bytes]]: + """ + _ageneric_api_call_with_fallbacks for anthropic_messages, with the + addition of mid-stream fallback handling (see + _aanthropic_messages_streaming_iterator). Parity with + _aresponses_with_streaming_fallbacks for the Responses API. + """ + from litellm.litellm_core_utils.core_helpers import safe_deep_copy + + # Snapshot the request kwargs before the primary attempt mutates them + # in place: _update_kwargs_with_deployment writes deployment-specific + # fields (deployment, model_info, api_base, tags, ...) into the + # SAME litellm_metadata/metadata dicts a shallow .copy() would still + # share, leaking primary-deployment metadata into the mid-stream + # fallback request. safe_deep_copy avoids deep-copying the full + # kwargs (which can hold non-deepcopyable logging handles/clients). + fallback_kwargs: Final[dict[str, object]] = kwargs.copy() # mutable-ok: mutated below before re-entry + if isinstance(fallback_kwargs.get("litellm_metadata"), dict): + fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"]) + if isinstance(fallback_kwargs.get("metadata"), dict): + fallback_kwargs["metadata"] = safe_deep_copy(fallback_kwargs["metadata"]) + fallback_kwargs["original_generic_function"] = original_function + + response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) + + if kwargs.get("stream") and hasattr(response, "__aiter__"): + return await self._aanthropic_messages_streaming_iterator( + response=cast("AsyncIterator[bytes]", response), # cast-ok: stream=True always returns a byte iterator + initial_kwargs=fallback_kwargs, + ) + return response + + async def _dispatch_generic_call_type( + self, + call_type: str, + original_function: Callable, + **kwargs: object, # kwargs-ok: forwarded verbatim to the per-call-type helper, shape varies per call site + ): + """ + factory_function's shared dispatch for call types with no + call-specific handling, except anthropic_messages: kept out of + factory_function's own async_wrapper (already at the repo's C901 + complexity ceiling) so routing its mid-stream fallback handling + (#24004) doesn't add another branch there. + """ + if call_type == "anthropic_messages": + return await self._aanthropic_messages_with_streaming_fallbacks( + original_function=original_function, **kwargs + ) + return await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) + def _generic_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs): """ Make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router @@ -5987,7 +6350,8 @@ class Router: "aget_skill", "adelete_skill", ): - return await self._ageneric_api_call_with_fallbacks( + return await self._dispatch_generic_call_type( + call_type=call_type, original_function=original_function, **kwargs, ) @@ -8341,6 +8705,7 @@ class Router: ) = litellm.get_llm_provider( model=deployment.litellm_params.model, custom_llm_provider=deployment.litellm_params.get("custom_llm_provider", None), + api_base=deployment.litellm_params.api_base, ) # done reading model["litellm_params"] # Check if provider is supported: either in enum or JSON-configured @@ -9448,6 +9813,8 @@ class Router: except Exception: model_info = None + deployment_is_mapped = deployment_is_catalog_mapped(model_info, model_info_dict) + # get llm provider litellm_model, llm_provider = "", "" try: @@ -9490,6 +9857,7 @@ class Router: "model_group": user_facing_model_group_name, "providers": [llm_provider], **model_info, + "supported_reasoning_efforts": None, } ) else: @@ -9567,6 +9935,11 @@ class Router: if model_info.get("rpm", None) is not None and _deployment_rpm is None: _deployment_rpm = model_info.get("rpm") + model_group_info.supported_reasoning_efforts = intersect_supported_reasoning_efforts( + model_group_info.supported_reasoning_efforts, + resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=deployment_is_mapped), + ) + if _deployment_tpm is not None: if total_tpm is None: total_tpm = 0 diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/router_utils/reasoning_effort_capability.py new file mode 100644 index 00000000000..3e4478e5e21 --- /dev/null +++ b/litellm/router_utils/reasoning_effort_capability.py @@ -0,0 +1,146 @@ +"""Resolve which reasoning_effort values a deployment, and by intersection a model group, accepts. + +The model map's supports_*_reasoning_effort flags are the only signal, and each level's polarity +mirrors how a request path reads that same flag. medium and high are unconditional for a reasoning +model. minimal and low are opt-out: openai/chat/gpt_5_transformation.py refuses them only when the +map says false. xhigh and max are opt-in. none is opt-out everywhere except the azure gpt-5 family, +whose config raises UnsupportedParamsError without an explicit true. + +xhigh is gated on the request path by the openai and azure gpt-5 configs. max is not gated there at +all: every entry carrying supports_max_reasoning_effort is Claude-family, and +anthropic/chat/transformation.py gates max on the output_config path while its reasoning_effort +path maps any level to a thinking budget. Making max opt-in is a deliberate trade, then, since an +explicit flag is the only signal that the tier is a real one rather than litellm rounding the level +to a budget, and a missing flag costs advisory metadata rather than a rejected request. + +A deployment the map describes with no effort flags at all resolves to None rather than to the +opt-out defaults. 689 of the map's 854 reasoning entries carry no flag, and the o-series, xai and +bedrock nova entries among them take neither none nor minimal, so composing a set out of the +defaults alone would advertise levels those providers reject. + +The advertisement order is the REASONING_EFFORT declaration order, which is presentation only. It +is not a strength scale and does not reconcile with bedrock's output_config ceiling order in +llms/bedrock/common_utils.py, which ranks max below xhigh while the thinking-budget constants rank +it above. +""" + +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final, get_args + +import litellm +from litellm.types.llms.openai import REASONING_EFFORT + +REASONING_EFFORT_ADVERTISEMENT_ORDER: Final = get_args(REASONING_EFFORT) +_EMPTY_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) + +_EFFORT_FLAGS: Final = ( + ("none", "supports_none_reasoning_effort"), + ("minimal", "supports_minimal_reasoning_effort"), + ("low", "supports_low_reasoning_effort"), + ("xhigh", "supports_xhigh_reasoning_effort"), + ("max", "supports_max_reasoning_effort"), +) +_OPT_OUT_EFFORTS: Final = ("minimal", "low") +_OPT_IN_EFFORTS: Final = ("xhigh", "max") +_UNCONDITIONAL_EFFORTS: Final = frozenset(("medium", "high")) + + +def _bare_model_entry(model_info: Mapping[str, object]) -> Mapping[str, object]: + """The unprefixed twin of a provider-prefixed map entry, which is where the flags often live: + azure/gpt-5-mini carries none of them while gpt-5-mini carries all three. The request-path + gates resolve through the same twin (_supports_factory, #20885), so reading it here is what + keeps the advertisement and the gate on the same answer.""" + key: Final = model_info.get("key") + provider: Final = model_info.get("litellm_provider") + if not isinstance(key, str) or not isinstance(provider, str) or not key.startswith(f"{provider}/"): + return _EMPTY_ENTRY + entry: Final[Mapping[str, object] | None] = litellm.model_cost.get(key.removeprefix(f"{provider}/")) + return entry if entry is not None else _EMPTY_ENTRY + + +def _declared_effort_flags(model_info: Mapping[str, object]) -> Mapping[str, object]: + bare: Final = _bare_model_entry(model_info) + return MappingProxyType( + { + effort: model_info.get(flag) if model_info.get(flag) is not None else bare.get(flag) + for effort, flag in _EFFORT_FLAGS + } + ) + + +def _supports_none_reasoning_effort(model_info: Mapping[str, object], flag: object) -> bool: + """Opt-in only where a request path refuses the level. AzureOpenAIGPT5Config raises + UnsupportedParamsError on reasoning_effort='none' without an explicit true, and it is selected + only for the gpt-5 family, so every other azure deployment keeps the opt-out default.""" + if model_info.get("litellm_provider") != "azure": + return flag is not False + + from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config + + key: Final = model_info.get("key") + if not isinstance(key, str) or not AzureOpenAIGPT5Config.is_model_gpt_5_model(key): + return flag is not False + return flag is True + + +def deployment_is_catalog_mapped( + resolved_model_info: Mapping[str, object] | None, + operator_model_info: Mapping[str, object], +) -> bool: + """Whether the model map described this deployment, as opposed to the operator describing it. + + Every deployment is registered in the cost map under its own id, so a mode the operator wrote + on an off-map deployment reads back here exactly like one the catalog supplied. Excluding it is + what stops such a deployment from claiming to be a known non-reasoning model and emptying the + levels its mapped siblings agree on. + """ + if resolved_model_info is None or resolved_model_info.get("mode") is None: + return False + return operator_model_info.get("mode") is None + + +def resolve_supported_reasoning_efforts( + model_info: Mapping[str, object], + *, + deployment_is_mapped: bool, +) -> tuple[str, ...] | None: + """None = nothing is known about this deployment, so it must not narrow its group; () = a known + model that accepts no effort level, which correctly empties the group. + + Telling those apart needs provenance the flattened ModelInfo does not carry. A deployment the + map does not describe arrives with supports_reasoning None, exactly like a mapped non-reasoning + model: 2273 of the map's 3165 entries omit the key rather than setting it false, so reading an + unset flag as () would let one custom deployment empty every level its mapped siblings agree + on. deployment_is_mapped is that provenance, and an operator who wants either answer for an + off-map deployment gets it by setting supports_reasoning explicitly. + """ + supports_reasoning: Final = model_info.get("supports_reasoning") + if supports_reasoning is not True: + return () if supports_reasoning is False or deployment_is_mapped else None + + flags: Final = _declared_effort_flags(model_info) + if all(value is None for value in flags.values()): + return None + + opt_out: Final = frozenset(effort for effort in _OPT_OUT_EFFORTS if flags[effort] is not False) + opt_in: Final = frozenset(effort for effort in _OPT_IN_EFFORTS if flags[effort] is True) + none_level: Final = ( + frozenset(("none",)) if _supports_none_reasoning_effort(model_info, flags["none"]) else frozenset() + ) + allowed: Final = opt_out | _UNCONDITIONAL_EFFORTS | opt_in | none_level + return tuple(effort for effort in REASONING_EFFORT_ADVERTISEMENT_ORDER if effort in allowed) + + +def intersect_supported_reasoning_efforts( + current: Sequence[str] | None, + resolved: Sequence[str] | None, +) -> tuple[str, ...] | None: + """Deployments without metadata (None) never narrow the group; an effort survives only when + every deployment with metadata accepts it, so the group offers nothing routing could reject.""" + if resolved is None: + return tuple(current) if current is not None else None + if current is None: + return tuple(resolved) + keep: Final = frozenset(current) & frozenset(resolved) + return tuple(effort for effort in REASONING_EFFORT_ADVERTISEMENT_ORDER if effort in keep) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index cc6eccbf3e0..d3b0f334163 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -685,6 +685,7 @@ ANTHROPIC_API_ONLY_HEADERS: Final = { # fails if calling anthropic on vertex ai class AnthropicThinkingParam(TypedDict, total=False): type: ReadOnly[Literal["enabled", "adaptive", "disabled"]] budget_tokens: int + display: ReadOnly[Literal["summarized", "omitted"]] class ANTHROPIC_HOSTED_TOOLS(str, Enum): diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 5665aa3277a..6ae2e31fe60 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1,4 +1,5 @@ import json +from collections.abc import Sequence from enum import Enum from typing import TYPE_CHECKING, Any, Final, Literal @@ -396,7 +397,7 @@ class OutputConfigBlock(TypedDict, total=False): class CommonRequestObject(TypedDict, total=False): # common request object across sync + async flows additionalModelRequestFields: dict - additionalModelResponseFieldPaths: list[str] + additionalModelResponseFieldPaths: Sequence[str] inferenceConfig: InferenceConfig system: list[SystemContentBlock] toolConfig: ToolConfigBlock diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index e7a3f825455..4a6c4a5bbb5 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1840,7 +1840,7 @@ ResponsesAPIStreamingResponse = Annotated[ ] -REASONING_EFFORT = Literal["none", "minimal", "low", "medium", "high", "xhigh"] +REASONING_EFFORT = Literal["none", "minimal", "low", "medium", "high", "xhigh", "max"] class OpenAIRealtimeStreamSession(TypedDict, total=False): diff --git a/litellm/types/router.py b/litellm/types/router.py index 9fd5cfa96ef..d4c735387a5 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -637,6 +637,7 @@ class ModelGroupInfo(BaseModel): supports_url_context: bool = Field(default=False) supports_reasoning: bool = Field(default=False) supports_function_calling: bool = Field(default=False) + supported_reasoning_efforts: tuple[str, ...] | None = Field(default=None) supported_openai_params: list[str] | None = Field(default=[]) configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None diff --git a/litellm/utils.py b/litellm/utils.py index 012e8785321..59a2de609d7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4130,7 +4130,7 @@ def get_optional_params( drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "together_ai": - optional_params = litellm.TogetherAIConfig().map_openai_params( + optional_params = litellm.TogetherAIChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, @@ -7898,7 +7898,7 @@ class ProviderConfigManager: LlmProviders.GALADRIEL: (lambda: litellm.GaladrielChatConfig(), False), LlmProviders.REPLICATE: (lambda: litellm.ReplicateConfig(), False), LlmProviders.HUGGINGFACE: (lambda: litellm.HuggingFaceChatConfig(), False), - LlmProviders.TOGETHER_AI: (lambda: litellm.TogetherAIConfig(), False), + LlmProviders.TOGETHER_AI: (lambda: litellm.TogetherAIChatConfig(), False), LlmProviders.OPENROUTER: (lambda: litellm.OpenrouterConfig(), False), LlmProviders.VERCEL_AI_GATEWAY: ( lambda: litellm.VercelAIGatewayConfig(), @@ -8610,6 +8610,12 @@ class ProviderConfigManager: ) return BedrockPassthroughConfig() + elif LlmProviders.BEDROCK_MANTLE == provider: + from litellm.llms.bedrock_mantle.passthrough.transformation import ( + BedrockMantlePassthroughConfig, + ) + + return BedrockMantlePassthroughConfig() elif LlmProviders.VLLM == provider or LlmProviders.HOSTED_VLLM == provider: from litellm.llms.vllm.passthrough.transformation import ( VLLMPassthroughConfig, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9f953e11df1..9ca8d9e1bac 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -37886,6 +37886,7 @@ "output_cost_per_token": 1e-07 }, "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo": { + "deprecation_date": "2026-02-06", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -37902,6 +37903,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-235B-A22B-Instruct-2507-tput": { + "deprecation_date": "2026-07-10", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "max_input_tokens": 262000, @@ -37914,6 +37916,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-235B-A22B-Thinking-2507": { + "deprecation_date": "2026-04-16", "input_cost_per_token": 6.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 256000, @@ -37926,6 +37929,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-235B-A22B-fp8-tput": { + "deprecation_date": "2026-02-06", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "max_input_tokens": 40000, @@ -37937,6 +37941,7 @@ "supports_tool_choice": false }, "together_ai/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { + "deprecation_date": "2026-06-04", "input_cost_per_token": 2e-06, "litellm_provider": "together_ai", "max_input_tokens": 256000, @@ -37949,11 +37954,15 @@ "supports_tool_choice": true }, "together_ai/deepseek-ai/DeepSeek-R1": { + "deprecation_date": "2026-05-14", "input_cost_per_token": 3e-06, "litellm_provider": "together_ai", "max_input_tokens": 128000, "max_output_tokens": 20480, "max_tokens": 20480, + "metadata": { + "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro" + }, "mode": "chat", "output_cost_per_token": 7e-06, "supports_function_calling": true, @@ -37962,6 +37971,7 @@ "supports_tool_choice": true }, "together_ai/deepseek-ai/DeepSeek-R1-0528-tput": { + "deprecation_date": "2026-02-03", "input_cost_per_token": 5.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 128000, @@ -37979,6 +37989,9 @@ "max_input_tokens": 65536, "max_output_tokens": 8192, "max_tokens": 8192, + "metadata": { + "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro" + }, "mode": "chat", "output_cost_per_token": 1.25e-06, "supports_function_calling": true, @@ -37987,9 +38000,13 @@ "supports_tool_choice": true }, "together_ai/deepseek-ai/DeepSeek-V3.1": { + "deprecation_date": "2026-05-14", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_tokens": 16384, + "metadata": { + "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro" + }, "mode": "chat", "output_cost_per_token": 1.7e-06, "source": "https://www.together.ai/models/deepseek-v3-1", @@ -38001,6 +38018,7 @@ "max_output_tokens": 16384 }, "together_ai/meta-llama/Llama-3.2-3B-Instruct-Turbo": { + "deprecation_date": "2026-03-06", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -38009,16 +38027,21 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo": { - "input_cost_per_token": 8.8e-07, + "input_cost_per_token": 1.04e-06, "litellm_provider": "together_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 8.8e-07, + "output_cost_per_token": 1.04e-06, + "source": "https://docs.together.ai/docs/serverless-models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true }, "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo-Free": { + "deprecation_date": "2025-11-13", "input_cost_per_token": 0, "litellm_provider": "together_ai", "mode": "chat", @@ -38029,6 +38052,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { + "deprecation_date": "2026-03-31", "input_cost_per_token": 2.7e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -38039,6 +38063,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Llama-4-Scout-17B-16E-Instruct": { + "deprecation_date": "2026-02-06", "input_cost_per_token": 1.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -38049,6 +38074,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": { + "deprecation_date": "2026-02-06", "input_cost_per_token": 3.5e-06, "litellm_provider": "together_ai", "mode": "chat", @@ -38059,6 +38085,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": { + "deprecation_date": "2026-02-25", "input_cost_per_token": 8.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -38069,6 +38096,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": { + "deprecation_date": "2026-03-06", "input_cost_per_token": 1.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -38079,6 +38107,7 @@ "supports_tool_choice": true }, "together_ai/mistralai/Mistral-7B-Instruct-v0.1": { + "deprecation_date": "2025-11-13", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -38087,6 +38116,7 @@ "supports_tool_choice": true }, "together_ai/mistralai/Mistral-Small-24B-Instruct-2501": { + "deprecation_date": "2026-04-02", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -38094,6 +38124,7 @@ "supports_tool_choice": true }, "together_ai/mistralai/Mixtral-8x7B-Instruct-v0.1": { + "deprecation_date": "2026-04-16", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -38106,6 +38137,9 @@ "together_ai/moonshotai/Kimi-K2-Instruct": { "input_cost_per_token": 1e-06, "litellm_provider": "together_ai", + "metadata": { + "successor": "together_ai/moonshotai/Kimi-K3" + }, "mode": "chat", "output_cost_per_token": 3e-06, "source": "https://www.together.ai/models/kimi-k2-instruct", @@ -38149,6 +38183,7 @@ "supports_tool_choice": true }, "together_ai/zai-org/GLM-4.5-Air-FP8": { + "deprecation_date": "2026-04-02", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "max_input_tokens": 128000, @@ -38166,6 +38201,9 @@ "max_input_tokens": 200000, "max_output_tokens": 200000, "max_tokens": 200000, + "metadata": { + "successor": "together_ai/zai-org/GLM-5.2" + }, "mode": "chat", "output_cost_per_token": 2.2e-06, "source": "https://www.together.ai/models/glm-4-6", @@ -38175,11 +38213,15 @@ "supports_tool_choice": true }, "together_ai/zai-org/GLM-4.7": { + "deprecation_date": "2026-04-02", "input_cost_per_token": 4.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 200000, "max_output_tokens": 200000, "max_tokens": 200000, + "metadata": { + "successor": "together_ai/zai-org/GLM-5.2" + }, "mode": "chat", "output_cost_per_token": 2e-06, "source": "https://www.together.ai/models/glm-4-7", @@ -38189,11 +38231,15 @@ "supports_tool_choice": true }, "together_ai/moonshotai/Kimi-K2.5": { + "deprecation_date": "2026-05-21", "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, + "metadata": { + "successor": "together_ai/moonshotai/Kimi-K3" + }, "mode": "chat", "output_cost_per_token": 2.8e-06, "source": "https://www.together.ai/models/kimi-k2-5", @@ -38203,9 +38249,13 @@ "supports_reasoning": true }, "together_ai/moonshotai/Kimi-K2-Instruct-0905": { + "deprecation_date": "2026-03-06", "input_cost_per_token": 1e-06, "litellm_provider": "together_ai", "max_input_tokens": 262144, + "metadata": { + "successor": "together_ai/moonshotai/Kimi-K3" + }, "mode": "chat", "output_cost_per_token": 3e-06, "source": "https://www.together.ai/models/kimi-k2-0905", @@ -38214,9 +38264,13 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-Next-80B-A3B-Instruct": { + "deprecation_date": "2026-04-02", "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, + "metadata": { + "successor": "together_ai/Qwen/Qwen3.7-Plus" + }, "mode": "chat", "output_cost_per_token": 1.5e-06, "source": "https://www.together.ai/models/qwen3-next-80b-a3b-instruct", @@ -38226,9 +38280,13 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-Next-80B-A3B-Thinking": { + "deprecation_date": "2026-02-25", "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, + "metadata": { + "successor": "together_ai/Qwen/Qwen3.6-Plus" + }, "mode": "chat", "output_cost_per_token": 1.5e-06, "source": "https://www.together.ai/models/qwen3-next-80b-a3b-thinking", @@ -38238,6 +38296,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3.5-397B-A17B": { + "deprecation_date": "2026-06-29", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, @@ -38249,6 +38308,292 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/MiniMaxAI/MiniMax-M3": { + "input_cost_per_token": 3e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/Prism-ML/Ternary-Bonsai-27B": { + "input_cost_per_token": 0.0, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/Qwen/Qwen3.5-9B": { + "input_cost_per_token": 1.7e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/Qwen/Qwen3.6-Plus": { + "input_cost_per_token": 5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_reasoning": true + }, + "together_ai/Qwen/Qwen3.7-Max": { + "input_cost_per_token": 1.25e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/Qwen/Qwen3.7-Plus": { + "input_cost_per_token": 3.2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.28e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/Qwen/Qwen3.8-2.4T-A95B": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1010000, + "max_output_tokens": 1010000, + "max_tokens": 1010000, + "mode": "chat", + "output_cost_per_token": 6.25e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/arize-ai/qwen-2-1.5b-instruct": { + "input_cost_per_token": 1e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/deepseek-ai/DeepSeek-V4-Flash-0731": { + "input_cost_per_token": 1.4e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/deepseek-ai/DeepSeek-V4-Pro": { + "input_cost_per_token": 1.74e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 512000, + "max_output_tokens": 512000, + "max_tokens": 512000, + "mode": "chat", + "output_cost_per_token": 3.48e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/deepseek-ai/DeepSeek-V4-Pro-0813": { + "input_cost_per_token": 1.32e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 3.96e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/google/gemma-3n-E4B-it": { + "input_cost_per_token": 6e-08, + "litellm_provider": "together_ai", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.2e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/google/gemma-4-31B-it": { + "input_cost_per_token": 3.9e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 9.7e-07, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/intfloat/multilingual-e5-large-instruct": { + "input_cost_per_token": 2e-08, + "litellm_provider": "together_ai", + "max_input_tokens": 514, + "max_tokens": 514, + "mode": "embedding", + "output_cost_per_token": 2e-08, + "output_vector_size": 1024, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/meta-llama/Llama-Guard-4-12B": { + "input_cost_per_token": 2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/meta-models/Muse-Glimmer-30B": { + "input_cost_per_token": 3.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/moonshotai/Kimi-K2.7-Code": { + "input_cost_per_token": 9.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/moonshotai/Kimi-K3": { + "input_cost_per_token": 3e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/nvidia/nemotron-3-ultra-550b-a55b": { + "input_cost_per_token": 6e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 512288, + "max_output_tokens": 512288, + "max_tokens": 512288, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/pearl-ai/gemma-4-31b-it": { + "input_cost_per_token": 2.8e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 8.6e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/thinkingmachines/Inkling": { + "input_cost_per_token": 1e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 4.05e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/thinkingmachines/Inkling-Small": { + "input_cost_per_token": 5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/zai-org/GLM-5.2": { + "input_cost_per_token": 1.4e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048575, + "max_output_tokens": 1048575, + "max_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "tts-1": { "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", @@ -49016,12 +49361,13 @@ "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -49048,12 +49394,13 @@ "output_cost_per_token": 1.32e-05, "output_cost_per_token_above_272k_tokens": 1.98e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -49080,12 +49427,13 @@ "output_cost_per_token": 1.32e-06, "output_cost_per_token_above_272k_tokens": 1.98e-06, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 03318718fb5..bbfbbbc1cfb 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3018 + "limit": 3016 }, "ANN002": { "limit": 71 @@ -9,7 +9,7 @@ "limit": 827 }, "ANN201": { - "limit": 2016 + "limit": 2015 }, "ANN202": { "limit": 852 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2919 + "limit": 2918 }, "C401": { "limit": 8 @@ -168,7 +168,7 @@ "limit": 3 }, "RET504": { - "limit": 176 + "limit": 175 }, "RUF012": { "limit": 240 diff --git a/tests/mcp_tests/conftest.py b/tests/mcp_tests/conftest.py index d1dc3ec7216..5823893afc0 100644 --- a/tests/mcp_tests/conftest.py +++ b/tests/mcp_tests/conftest.py @@ -7,6 +7,7 @@ import pytest import litellm import asyncio +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @pytest.fixture(scope="session") @@ -38,6 +39,8 @@ def setup_and_teardown(): yield # Teardown code (executes after the yield point) + # LoggingWorker carries still-queued coroutines onto the next test's loop, where they'd log into that test's callbacks + asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) loop.close() # Close the loop created earlier asyncio.set_event_loop(None) # Remove the reference to the loop diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py index 1903f29001f..fc9f675f837 100644 --- a/tests/mcp_tests/test_mcp_logging.py +++ b/tests/mcp_tests/test_mcp_logging.py @@ -1,6 +1,9 @@ import os import pytest import asyncio +import subprocess +import sys +from pathlib import Path from typing import Optional from unittest.mock import AsyncMock, patch @@ -458,3 +461,45 @@ async def test_mcp_tool_call_hook(): logged_standard_logging_payload is not None ), "Standard logging payload should not be None" assert logged_standard_logging_payload["response_cost"] == 1.42 + + +_QUEUED_LOGGING_OUTLIVES_TEST = ''' +import time + +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + +ran_at = [] + + +async def _record_run(): + ran_at.append(time.monotonic()) + + +async def test_1_leaves_logging_queued_behind_a_stopped_worker(): + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(_record_run()) + await GLOBAL_LOGGING_WORKER.stop() + assert ran_at == [] + + +async def test_2_starts_after_the_previous_tests_logging_ran(): + started_at = time.monotonic() + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(_record_run()) + await GLOBAL_LOGGING_WORKER.flush() + assert [t < started_at for t in ran_at] == [True, False] +''' + + +def test_logging_queued_by_one_test_is_drained_before_the_next(tmp_path: Path): + """Regression: a logging coroutine queued by one test must not run inside a later test (it would log into that + test's callbacks, which is how test_mcp_tool_call_hook captured a gpt-4o-mini payload under xdist).""" + (tmp_path / "conftest.py").write_text((Path(__file__).parent / "conftest.py").read_text()) + (tmp_path / "pyproject.toml").write_text('[tool.pytest.ini_options]\nasyncio_mode = "auto"\n') + (tmp_path / "test_queued_logging.py").write_text(_QUEUED_LOGGING_OUTLIVES_TEST) + result = subprocess.run( + [sys.executable, "-m", "pytest", "-q", "-p", "no:cacheprovider", "test_queued_logging.py"], + cwd=tmp_path, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py b/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py new file mode 100644 index 00000000000..30544a8bb81 --- /dev/null +++ b/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py @@ -0,0 +1,307 @@ +""" +Real-Postgres coverage for the /team/member_add vs /team/delete race (LIT-5544), and for +/team/member_delete's participation in the same lock. + +A member_add that validated the team before a delete began could previously still commit +its writes after the delete's reference sweeps had already run, leaving a user record and +a membership row pointing at a team id that no longer exists. Neither side of that race can +be forced by a sequential script: it needs one request to be genuinely mid-flight while the +other commits. A mocked prisma cannot arbitrate that either, since the property under test +is whether Postgres's own advisory lock actually serializes the two requests. + +These tests pin the interleaving the same way test_access_group_team_sync.py does: a second +real connection holds the team's advisory lock in its own transaction, so the function under +test is provably blocked on it rather than hoping a sleep lands in the right gap. +""" + +import asyncio +import json +import os +from contextlib import asynccontextmanager +from datetime import timedelta +from unittest.mock import MagicMock + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import ( + DeleteTeamRequest, + LitellmUserRoles, + Member, + TeamMemberAddRequest, + UserAPIKeyAuth, +) +from litellm.caching.caching import DualCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + +TEAM = "lit5544-race-team" +USER = "lit5544-race-user" +_DELETE_SEEDED = 'DELETE FROM "LiteLLM_TeamMembership" WHERE team_id = $1' +_DELETE_USER = 'DELETE FROM "LiteLLM_UserTable" WHERE user_id = $1' +_DELETE_TEAM = 'DELETE FROM "LiteLLM_TeamTable" WHERE team_id = $1' +_LOCK_SQL = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked" + + +@asynccontextmanager +async def _clean_db(): + """Connects inside the running test's loop: an async fixture would be torn up on a + different loop than the test body, which prisma's engine lock refuses outright.""" + from prisma import Prisma + + if not os.getenv("DATABASE_URL"): + pytest.fail("DATABASE_URL is required; these tests must not silently skip") + + db = Prisma() + await db.connect() + try: + await db.execute_raw(_DELETE_SEEDED, TEAM) + await db.execute_raw(_DELETE_USER, USER) + await db.execute_raw(_DELETE_TEAM, TEAM) + yield db + finally: + await db.execute_raw(_DELETE_SEEDED, TEAM) + await db.execute_raw(_DELETE_USER, USER) + await db.execute_raw(_DELETE_TEAM, TEAM) + await db.disconnect() + + +@asynccontextmanager +async def _real_prisma_client(): + """The full app-level PrismaClient, not the raw generated client: add_new_member reads + and writes through PrismaClient.get_data/insert_data, which the raw client doesn't have.""" + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + client = PrismaClient(database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj) + await client.connect() + try: + yield client + finally: + await client.db.disconnect() + + +def _admin_auth(): + return UserAPIKeyAuth(user_id="lit5544-admin", api_key="sk-lit5544", user_role=LitellmUserRoles.PROXY_ADMIN.value) + + +@pytest.mark.asyncio +async def test_member_add_blocked_by_delete_writes_no_dangling_reference(): + """ + member_add re-reads the team under the advisory lock before writing anything. When a + delete already holds that lock and then removes the row, member_add's re-read must see + the row gone and raise, without ever calling the write that appends the user/membership + references, which is the only way this leaves zero trace after the delete wins. + """ + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _add_team_members_to_team, + ) + + async with _clean_db() as db: + await db.litellm_teamtable.create(data={"team_id": TEAM, "team_alias": TEAM, "members_with_roles": "[]"}) + + async with _real_prisma_client() as prisma_client: + from prisma import Prisma + + blocker = Prisma() + await blocker.connect() + lock_acquired = asyncio.Event() + + async def add_member(): + lock_acquired.set() + await _add_team_members_to_team( + data=TeamMemberAddRequest( + team_id=TEAM, + member=Member(user_id=USER, role="user"), + max_budget_in_team=5.0, + ), + complete_team_data=LiteLLM_TeamTable(team_id=TEAM, members_with_roles=[]), + prisma_client=prisma_client, + user_api_key_dict=_admin_auth(), + litellm_proxy_admin_name="lit5544-admin", + ) + + try: + async with blocker.tx(timeout=timedelta(seconds=30)) as held: + await held.query_raw(_LOCK_SQL, TEAM) + task = asyncio.create_task(add_member()) + await lock_acquired.wait() + await asyncio.sleep(0.2) + assert not task.done(), "member_add did not wait on the team's advisory lock" + + # the delete wins the race: strip the team row while the lock is held + await held.execute_raw(_DELETE_TEAM, TEAM) + + with pytest.raises(HTTPException) as exc_info: + await asyncio.wait_for(task, timeout=30) + assert exc_info.value.status_code == 404 + finally: + await blocker.disconnect() + + user_row = await db.litellm_usertable.find_unique(where={"user_id": USER}) + assert user_row is None, "member_add must not have written a user row for a team that was gone under its lock" + + membership_row = await db.litellm_teammembership.find_first(where={"team_id": TEAM, "user_id": USER}) + assert membership_row is None + + +@pytest.mark.asyncio +async def test_member_delete_blocked_by_member_add_removes_from_the_fresh_roster(): + """ + team_member_delete takes the same advisory lock and re-reads the roster under it, so a + member_add that committed while member_delete was waiting on the lock is not silently + undone. Without the re-read, member_delete would compute its new roster from the stale + snapshot it validated against before the lock, and its write would overwrite the + member_add's addition right back out even though member_add's request already succeeded. + """ + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import TeamMemberDeleteRequest + from litellm.proxy.management_endpoints.team_endpoints import team_member_delete + + other_user = f"{USER}-other" + seeded_roster = '[{"user_id": "%s", "user_email": null, "role": "user"}]' % USER + winning_add_roster = ( + '[{"user_id": "%s", "user_email": null, "role": "user"}, ' + '{"user_id": "%s", "user_email": null, "role": "user"}]' % (USER, other_user) + ) + + async with _clean_db() as db: + await db.litellm_teamtable.create( + data={"team_id": TEAM, "team_alias": TEAM, "members_with_roles": seeded_roster} + ) + + async with _real_prisma_client() as prisma_client: + original_prisma_client = proxy_server_module.prisma_client + proxy_server_module.prisma_client = prisma_client + + try: + from prisma import Prisma + + blocker = Prisma() + await blocker.connect() + lock_acquired = asyncio.Event() + + async def run_delete(): + lock_acquired.set() + return await team_member_delete( + data=TeamMemberDeleteRequest(team_id=TEAM, user_id=USER), + user_api_key_dict=_admin_auth(), + ) + + try: + async with blocker.tx(timeout=timedelta(seconds=30)) as held: + await held.query_raw(_LOCK_SQL, TEAM) + task = asyncio.create_task(run_delete()) + await lock_acquired.wait() + await asyncio.sleep(0.2) + assert not task.done(), "member_delete did not wait on the team's advisory lock" + + # member_add wins the race: it adds `other_user` while holding the lock + await held.litellm_teamtable.update( + where={"team_id": TEAM}, + data={"members_with_roles": winning_add_roster}, + ) + + await asyncio.wait_for(task, timeout=30) + finally: + await blocker.disconnect() + finally: + proxy_server_module.prisma_client = original_prisma_client + + team_row = await db.litellm_teamtable.find_unique(where={"team_id": TEAM}) + raw_roster = team_row.members_with_roles + parsed_roster = json.loads(raw_roster) if isinstance(raw_roster, str) else raw_roster + remaining_ids = {m["user_id"] for m in parsed_roster} + assert remaining_ids == {other_user}, ( + "member_delete must remove only the user it targeted from the roster it actually " + "committed to, not silently drop the member the winning add just committed" + ) + + +@pytest.mark.asyncio +async def test_delete_blocked_by_member_add_sweeps_the_fresh_reference(): + """ + A member_add that wins the lock race writes its reference and releases the lock; the + delete that was waiting on it must then run its locked sweep against the row as it + actually is, not a stale snapshot, and reap that reference rather than leaving it + stranded on a team id the delete is about to remove. + """ + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.management_endpoints.team_endpoints import delete_team + + async with _clean_db() as db: + await db.litellm_teamtable.create(data={"team_id": TEAM, "team_alias": TEAM, "members_with_roles": "[]"}) + + async with _real_prisma_client() as prisma_client: + proxy_logging_obj = prisma_client.proxy_logging_obj + original_prisma_client = proxy_server_module.prisma_client + original_admin_name = proxy_server_module.litellm_proxy_admin_name + original_proxy_logging_obj = proxy_server_module.proxy_logging_obj + original_cache = proxy_server_module.user_api_key_cache + original_router = proxy_server_module.llm_router + proxy_server_module.prisma_client = prisma_client + proxy_server_module.litellm_proxy_admin_name = "lit5544-admin" + proxy_server_module.proxy_logging_obj = proxy_logging_obj + proxy_server_module.user_api_key_cache = original_cache or proxy_logging_obj.internal_usage_cache + proxy_server_module.llm_router = None + + async def restore(): + proxy_server_module.prisma_client = original_prisma_client + proxy_server_module.litellm_proxy_admin_name = original_admin_name + proxy_server_module.proxy_logging_obj = original_proxy_logging_obj + proxy_server_module.user_api_key_cache = original_cache + proxy_server_module.llm_router = original_router + + try: + from prisma import Prisma + + blocker = Prisma() + await blocker.connect() + lock_acquired = asyncio.Event() + + async def run_delete(): + lock_acquired.set() + return await delete_team( + data=DeleteTeamRequest(team_ids=[TEAM]), + http_request=MagicMock(), + user_api_key_dict=_admin_auth(), + litellm_changed_by="lit5544-admin", + ) + + try: + async with blocker.tx(timeout=timedelta(seconds=30)) as held: + await held.query_raw(_LOCK_SQL, TEAM) + task = asyncio.create_task(run_delete()) + await lock_acquired.wait() + await asyncio.sleep(0.3) + assert not task.done(), "delete_team did not wait on the team's advisory lock" + + # member_add wins the race: write the reference while holding the lock + await held.litellm_usertable.upsert( + where={"user_id": USER}, + data={ + "create": {"user_id": USER, "teams": [TEAM]}, + "update": {"teams": {"push": [TEAM]}}, + }, + ) + await held.litellm_teammembership.create(data={"team_id": TEAM, "user_id": USER}) + await held.litellm_teamtable.update( + where={"team_id": TEAM}, + data={"members_with_roles": '[{"user_id": "%s", "role": "user"}]' % USER}, + ) + + await asyncio.wait_for(task, timeout=30) + finally: + await blocker.disconnect() + finally: + await restore() + + team_row = await db.litellm_teamtable.find_unique(where={"team_id": TEAM}) + assert team_row is None + + user_row = await db.litellm_usertable.find_unique(where={"user_id": USER}) + assert user_row is not None and TEAM not in user_row.teams, ( + "delete_team's locked sweep must reap the reference member_add wrote just before losing the lock" + ) + + membership_row = await db.litellm_teammembership.find_first(where={"team_id": TEAM, "user_id": USER}) + assert membership_row is None diff --git a/tests/proxy_behavior/management/test_team_member_reset_spend.py b/tests/proxy_behavior/management/test_team_member_reset_spend.py new file mode 100644 index 00000000000..ec2c78139fe --- /dev/null +++ b/tests/proxy_behavior/management/test_team_member_reset_spend.py @@ -0,0 +1,152 @@ +import uuid + +import pytest + +from .actors import Actor +from .conftest import create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + +_SEED_SPEND = 5.0 +_RESET_TO = 2.0 + + +# POST /team/{team_id}/member/{user_id}/reset_spend. The handler gate is +# _verify_team_access (proxy admin / team admin of this team / org admin of +# the team's org) — the same gate /team/member_update uses, so this mirrors +# that file's matrix exactly. +_MATRIX = [ + ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), + ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403), + ("alpha/owner", Actor.OWNER, "alpha", 403), + ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403), + ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403), + ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 403), + ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403), + ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), +] + + +async def _seed_target(prisma, world, shape: str, team_id: str, member_id: str) -> None: + if shape == "alpha": + await create_scratch_team( + prisma, + team_id, + organization_id=world.org_a_id, + admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id], + ) + elif shape == "beta": + await create_scratch_team(prisma, team_id, organization_id=world.org_b_id) + else: # pragma: no cover - guard + pytest.fail(f"unknown shape={shape}") + await prisma.db.litellm_teammembership.create( + data={"user_id": member_id, "team_id": team_id, "spend": _SEED_SPEND} + ) + + +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_member_reset_spend_authz_matrix( + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + member_id = scratch.tag("member") + await _seed_target(prisma, world, shape, scratch.prefix, member_id) + caller = world.keys[actor] + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_spend", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"reset_to": _RESET_TO}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teammembership.find_unique( + where={"user_id_team_id": {"user_id": member_id, "team_id": scratch.prefix}} + ) + assert row is not None + if expected_status == 200: + assert row.spend == _RESET_TO + else: + assert row.spend == _SEED_SPEND, "denied but spend reset" + + +async def test_team_member_reset_spend_missing_team_is_404(proxy_client, world): + resp = await proxy_client.post( + f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_spend", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"reset_to": 0.0}, + ) + assert resp.status_code == 404, resp.text + + +async def test_team_member_reset_spend_missing_membership_is_404( + proxy_client, prisma, scratch, world +): + """A well-formed team but a user_id with no LiteLLM_TeamMembership row is 404.""" + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{uuid.uuid4().hex}/reset_spend", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"reset_to": 0.0}, + ) + assert resp.status_code == 404, resp.text + + +async def test_team_member_reset_spend_above_current_spend_is_400( + proxy_client, prisma, scratch, world +): + member_id = scratch.tag("member") + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + await prisma.db.litellm_teammembership.create( + data={"user_id": member_id, "team_id": scratch.prefix, "spend": 1.0} + ) + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_spend", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"reset_to": 5.0}, + ) + assert resp.status_code == 400, resp.text + + +async def test_team_member_reset_spend_team_admin_cannot_reset_own_spend( + proxy_client, prisma, scratch, world +): + """A team admin targeting their own LiteLLM_TeamMembership row is 403: unchecked, an + admin could repeatedly zero their own spend right before it crosses their per-member + cap, consuming the shared team budget without the configured limit ever binding.""" + team_admin = world.keys[Actor.TEAM_ADMIN] + await create_scratch_team( + prisma, + scratch.prefix, + organization_id=world.org_a_id, + admin_user_ids=[team_admin.user_id], + ) + await prisma.db.litellm_teammembership.create( + data={"user_id": team_admin.user_id, "team_id": scratch.prefix, "spend": _SEED_SPEND} + ) + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{team_admin.user_id}/reset_spend", + headers={"Authorization": f"Bearer {team_admin.cleartext}"}, + json={"reset_to": 0.0}, + ) + assert resp.status_code == 403, resp.text + row = await prisma.db.litellm_teammembership.find_unique( + where={"user_id_team_id": {"user_id": team_admin.user_id, "team_id": scratch.prefix}} + ) + assert row is not None and row.spend == _SEED_SPEND, "denied but spend reset" diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 21dbf3e090f..ceaabf0a70f 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -1169,6 +1169,22 @@ async def test_create_user_default_budget(prisma_client, user_role): # noqa: F8 assert mock_client.call_args.kwargs["data"]["budget_duration"] is None +def _member_add_tx_cm(team_table): + """Transaction whose member writes land on whatever tables are mocked on `prisma_client.db`""" + + class _Tx: + query_raw = AsyncMock(return_value=[{"members_with_roles": []}]) + litellm_teamtable = team_table + + def __getattr__(self, table_name): + return getattr(litellm.proxy.proxy_server.prisma_client.db, table_name) + + tx_cm = MagicMock() + tx_cm.__aenter__ = AsyncMock(return_value=_Tx()) + tx_cm.__aexit__ = AsyncMock(return_value=None) + return tx_cm + + @pytest.mark.parametrize("new_member_method", ["user_id", "user_email"]) @pytest.mark.asyncio @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") @@ -1230,7 +1246,7 @@ async def test_create_team_member_add(prisma_client, new_member_method): # noqa ) ) mock_litellm_usertable.upsert = mock_client - mock_litellm_usertable.find_many = AsyncMock(return_value=None) + mock_litellm_usertable.find_many = AsyncMock(return_value=[]) # Mock find_first for user_email validation (returns None for new users) mock_litellm_usertable.find_first = AsyncMock(return_value=None) # Mock find_unique for user_id validation (returns None for new users) @@ -1245,12 +1261,7 @@ async def test_create_team_member_add(prisma_client, new_member_method): # noqa return_value=LiteLLM_TeamTableCachedObj(team_id="1234") ) - tx_mock = AsyncMock() - tx_mock.query_raw = AsyncMock(return_value=[{"members_with_roles": []}]) - tx_mock.litellm_teamtable = team_mock_client - tx_cm = MagicMock() - tx_cm.__aenter__ = AsyncMock(return_value=tx_mock) - tx_cm.__aexit__ = AsyncMock(return_value=None) + tx_cm = _member_add_tx_cm(team_mock_client) original_tx = litellm.proxy.proxy_server.prisma_client.tx litellm.proxy.proxy_server.prisma_client.tx = MagicMock( return_value=tx_cm @@ -1432,7 +1443,7 @@ async def test_create_team_member_add_team_admin( ) ) mock_litellm_usertable.upsert = mock_client - mock_litellm_usertable.find_many = AsyncMock(return_value=None) + mock_litellm_usertable.find_many = AsyncMock(return_value=[]) # Mock find_first for user_email validation (returns None for new users) mock_litellm_usertable.find_first = AsyncMock(return_value=None) # Mock find_unique for user_id validation (returns None for new users) @@ -1443,12 +1454,7 @@ async def test_create_team_member_add_team_admin( return_value=LiteLLM_TeamTableCachedObj(team_id="1234") ) - tx_mock = AsyncMock() - tx_mock.query_raw = AsyncMock(return_value=[{"members_with_roles": []}]) - tx_mock.litellm_teamtable = team_mock_client - tx_cm = MagicMock() - tx_cm.__aenter__ = AsyncMock(return_value=tx_mock) - tx_cm.__aexit__ = AsyncMock(return_value=None) + tx_cm = _member_add_tx_cm(team_mock_client) with ( patch.object( diff --git a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py b/tests/test_litellm/caching/test_redis_cluster_node_isolation.py index f4cd3ab20ef..c16ceec8c31 100644 --- a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py +++ b/tests/test_litellm/caching/test_redis_cluster_node_isolation.py @@ -28,6 +28,14 @@ if TYPE_CHECKING: from redis.asyncio.cluster import RedisCluster as _AsyncRedisClusterType +class _NodeClassWithPerConnectionRecovery: + def update_active_connections_for_reconnect(self) -> None: ... + + +class _NodeClassWithoutPerConnectionRecovery: + pass + + class _FakeClusterNode: def __init__(self, name: str, raises: Exception | None = None, response: object = None) -> None: self.name = name @@ -47,7 +55,9 @@ class _FakeNodesManager: def _build_cluster_instance() -> "_AsyncRedisClusterType": - cluster_cls = get_litellm_async_redis_cluster_class() + cluster_cls = get_litellm_async_redis_cluster_class( + cluster_node_class=_NodeClassWithoutPerConnectionRecovery + ) instance = cluster_cls.__new__(cluster_cls) instance.RedisClusterRequestTTL = 1 instance.reinitialize_counter = 0 @@ -58,6 +68,33 @@ def _build_cluster_instance() -> "_AsyncRedisClusterType": return instance +def test_per_connection_recovery_redis_py_gets_the_unmodified_upstream_class() -> None: + """Regression (redis-py 8.x): when upstream ClusterNode already recovers a node-level + connection error per-connection, the factory must NOT install the copied override, + whose node.disconnect() also kills connections other coroutines are mid-operation on.""" + from redis.asyncio.cluster import RedisCluster + + cluster_cls = get_litellm_async_redis_cluster_class( + cluster_node_class=_NodeClassWithPerConnectionRecovery + ) + + assert cluster_cls is RedisCluster + + +def test_pre_recovery_redis_py_still_gets_the_node_isolation_override() -> None: + """Old redis-py (5.x) responds to a node-level error with a full-cluster aclose(), + so those versions must keep litellm's per-node isolation override.""" + from redis.asyncio.cluster import RedisCluster + + cluster_cls = get_litellm_async_redis_cluster_class( + cluster_node_class=_NodeClassWithoutPerConnectionRecovery + ) + + assert cluster_cls is not RedisCluster + assert issubclass(cluster_cls, RedisCluster) + assert "_execute_command" in cluster_cls.__dict__ + + @pytest.mark.asyncio @pytest.mark.parametrize("error_cls", [RedisConnectionError, RedisTimeoutError]) async def test_node_level_error_resets_only_that_node_not_the_whole_client(error_cls: type[Exception]) -> None: 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 1fb74b2b7bf..6ca48ce63b8 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 @@ -2,7 +2,7 @@ import datetime import json import os import unittest -from typing import TYPE_CHECKING, List, Literal, Optional, Tuple +from typing import TYPE_CHECKING, Final, List, Literal, Optional, Tuple from unittest.mock import ANY, MagicMock, Mock, patch import httpx @@ -1585,10 +1585,16 @@ def test_map_reasoning_effort_adds_summary_detailed(monkeypatch): assert result_dict["summary"] == "custom_summary" print("✓ Dict input is passed through without modification") - # Test 5: None/unknown values return None - result_unknown = handler._map_reasoning_effort("unknown_value") - assert result_unknown is None - print("✓ Unknown reasoning_effort values return None") + # Test 5: every REASONING_EFFORT level reaches the provider, and anything else (a typo, an + # unshipped level, "default") is dropped so the request still succeeds at the provider default + from litellm.types.llms.openai import Reasoning + + for effort in ("max", "xhigh", "none"): + result_passthrough = handler._map_reasoning_effort(effort) + assert result_passthrough == Reasoning(effort=effort) + for dropped in ("ultra", "hgih", "unknown_value", "", "default"): + assert handler._map_reasoning_effort(dropped) is None + print("✓ Enumerated levels pass through and unknown ones are dropped") print( "✓ All reasoning_effort behaviors work correctly with flag/env var control" @@ -2438,6 +2444,32 @@ def test_map_optional_params_preserves_reasoning_summary(): assert responses_api_request["reasoning"]["summary"] == "detailed" +@pytest.mark.parametrize("reasoning_effort", ["max", "high"]) +def test_transform_request_bedrock_mantle_tools_keeps_reasoning_effort(monkeypatch, reasoning_effort): + """Regression for reasoning_effort=max being dropped on the chat -> Responses bridge (issue #38084).""" + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + monkeypatch.setattr(litellm, "reasoning_auto_summary", False) + monkeypatch.delenv("LITELLM_REASONING_AUTO_SUMMARY", raising=False) + handler: Final = LiteLLMResponsesTransformationHandler() + + result: Final = handler.transform_request( + model="openai.gpt-5.6-sol", + messages=[{"role": "user", "content": "Say pong"}], + optional_params={ + "reasoning_effort": reasoning_effort, + "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + }, + litellm_params={"custom_llm_provider": "bedrock_mantle"}, + headers={}, + litellm_logging_obj=Mock(), + ) + + assert result["reasoning"] == {"effort": reasoning_effort} + + def test_map_optional_params_tool_choice_chat_nested_to_responses_api(): """Chat tool_choice must become Responses ToolChoiceFunction (top-level name).""" from litellm.completion_extras.litellm_responses_transformation.transformation import ( 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 c8c36032793..6f513ce1bd4 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 @@ -478,10 +478,10 @@ def test_generic_cost_per_token_minimax_m3_above_512k_tokens(_local_model_cost_m ], ) def test_generic_cost_per_token_bedrock_mantle_gpt56_long_context(_local_model_cost_map, model): - """Bedrock GPT-5.6 supports a 1M context window, billed at the long-context rates above 272K.""" + """Bedrock GPT-5.6 enforces a 1,050,000-token context window, billed at the long-context rates above 272K.""" model_cost_map = litellm.model_cost[model] - assert model_cost_map["max_input_tokens"] == 1000000 + assert model_cost_map["max_input_tokens"] == 1050000 cached_tokens = 100000 completion_tokens = 1000 diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index aa7f2990c13..33a33711698 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -255,3 +255,26 @@ class TestRedactNestedMatchAndRegexKeys: def test_passes_through_none_and_str(self): assert redact_nested_match_and_regex_keys(None) is None assert redact_nested_match_and_regex_keys("plain") == "plain" + + +class TestIsExpectedClientError: + def test_status_ranges(self): + from litellm.litellm_core_utils.core_helpers import is_expected_client_error + + class WithStatusCode(Exception): + def __init__(self, status_code): + self.status_code = status_code + + class WithCode(Exception): + def __init__(self, code): + self.code = code + + assert is_expected_client_error(WithStatusCode(400)) is True + assert is_expected_client_error(WithStatusCode(429)) is True + assert is_expected_client_error(WithStatusCode(499)) is True + assert is_expected_client_error(WithStatusCode(500)) is False + assert is_expected_client_error(WithStatusCode(399)) is False + assert is_expected_client_error(WithCode("403")) is True + assert is_expected_client_error(WithCode("invalid_request_error")) is False + assert is_expected_client_error(Exception("no status")) is False + assert is_expected_client_error(None) is False diff --git a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py index bda7ab4afc6..6cacd119030 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py +++ b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py @@ -133,3 +133,54 @@ class TestGetLlmProviderRejectsAttackerSmuggledApiBase: assert provider == "groq" assert dynamic_api_key == "server-real-groq-key" + + +class TestTogetherApiBaseResolvesProvider: + """ + Regression for the Together host migration: both the current + ``api.together.ai`` host and the legacy ``api.together.xyz`` host must + resolve to ``together_ai`` when passed as ``api_base``. Before the fix + the endpoint list carried the legacy host but the provider-mapping + chain had no branch for it, so the match fell through with a None + provider and the deployment failed with "LLM Provider NOT provided". + """ + + @pytest.mark.parametrize( + "api_base", + [ + "https://api.together.ai/v1", + "https://api.together.xyz/v1", + ], + ) + def test_together_api_base_resolves_to_together_ai(self, api_base, monkeypatch): + monkeypatch.setenv("TOGETHER_API_KEY", "together-key-from-env") + + model, provider, dynamic_api_key, returned_api_base = get_llm_provider( + model="some-model", + api_base=api_base, + ) + + assert provider == "together_ai" + assert dynamic_api_key == "together-key-from-env" + assert returned_api_base == api_base + assert model == "some-model" + + def test_explicit_api_key_beats_together_env_key(self, monkeypatch): + monkeypatch.setenv("TOGETHER_API_KEY", "together-key-from-env") + + _, provider, dynamic_api_key, _ = get_llm_provider( + model="some-model", + api_base="https://api.together.ai/v1", + api_key="explicit-caller-key", + ) + + assert provider == "together_ai" + assert dynamic_api_key == "explicit-caller-key" + + def test_together_default_api_base_is_together_ai(self, monkeypatch): + monkeypatch.delenv("TOGETHER_AI_API_BASE", raising=False) + + _, provider, _, api_base = get_llm_provider(model="together_ai/some-model") + + assert provider == "together_ai" + assert api_base == "https://api.together.ai/v1" diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 1cac0cce4ef..f93daa61570 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -6,7 +6,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest - import time import httpx @@ -3867,6 +3866,90 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob assert payload["litellm_call_id"] == call_id +# ── Azure Model Router selected-model attribution ──────────────────────────── + + +def _model_router_response(selected_model: str, stamp: bool): + """A ModelResponse as AzureModelRouterConfig hands it back, with or without the stamp.""" + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + ) + from litellm.types.utils import ModelResponse + + response = ModelResponse(model=selected_model) + response._hidden_params = ( + {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {} + ) + return response + + +def test_standard_logging_payload_uses_stamped_model_router_model(logging_obj): + """ + The selected model must win off the stamp, not off "model-router" appearing in the + requested model. An operator whose model group is named anything else was invisible + to the name check, so their logs and spend rows named the router instead. + """ + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.datetime.now() + payload = get_standard_logging_object_payload( + kwargs={ + "model": "azure_ai/smart-pick", + "custom_llm_provider": "azure_ai", + "messages": [], + "litellm_params": {"metadata": {}}, + }, + init_response_obj=_model_router_response( + "azure_ai/grok-4-1-fast-reasoning", stamp=True + ), + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["model"] == "azure_ai/grok-4-1-fast-reasoning" + + +def test_standard_logging_payload_keeps_requested_model_without_router_stamp( + logging_obj, +): + """ + Control for the test above: an ordinary azure_ai deployment is unaffected, so the stamp + is what redirects attribution rather than the response model winning unconditionally. + """ + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.datetime.now() + payload = get_standard_logging_object_payload( + kwargs={ + "model": "azure_ai/smart-pick", + "custom_llm_provider": "azure_ai", + "messages": [], + "litellm_params": {"metadata": {}}, + }, + init_response_obj=_model_router_response( + "azure_ai/grok-4-1-fast-reasoning", stamp=False + ), + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["model"] == "azure_ai/smart-pick" + + def _make_dict_logging_obj(): """Build a Logging instance configured for a non-streaming dict result.""" obj = LitellmLogging( @@ -5595,3 +5678,62 @@ def test_get_custom_logger_compatible_class_finds_v2_newrelic(monkeypatch): logging_module._in_memory_loggers.clear() monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) is_otel_v2_enabled.cache_clear() + + +class _ClientError(Exception): + def __init__(self, status_code, message): + self.status_code = status_code + self.message = message + super().__init__(message) + + +def _raise_and_catch(exc): + try: + raise exc + except Exception as caught: + return caught + + +def test_get_error_information_skips_traceback_for_expected_4xx(monkeypatch): + """Regression for LIT-6043: expected client (4xx) errors must not pay for + traceback.format_tb on every rejected request unless + litellm.log_client_error_tracebacks is enabled.""" + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + client_exc = _raise_and_catch(_ClientError(status_code=403, message="team does not allow model")) + assert client_exc.__traceback__ is not None + result = StandardLoggingPayloadSetup.get_error_information(client_exc) + assert result["traceback"] == "" + + server_exc = _raise_and_catch(_ClientError(status_code=500, message="boom")) + result = StandardLoggingPayloadSetup.get_error_information(server_exc) + assert "test_litellm_logging" in result["traceback"] + + monkeypatch.setattr(litellm, "log_client_error_tracebacks", True) + result = StandardLoggingPayloadSetup.get_error_information(client_exc) + assert "test_litellm_logging" in result["traceback"] + + +def test_failure_handler_helper_fn_builds_payload_once_per_exception(): + """Regression for LIT-6043: async and sync failure handlers both call + _failure_handler_helper_fn for the same failed request; the standardized + payload must be built once, not once per handler.""" + obj = LitellmLogging( + model="gpt-4o", + messages=[{"role": "user", "content": "Hey"}], + stream=False, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="lit-6043-1", + function_id="f", + ) + exc = _raise_and_catch(_ClientError(status_code=400, message="invalid model")) + obj._failure_handler_helper_fn(exception=exc, traceback_exception="") + first_payload = obj.model_call_details["standard_logging_object"] + assert first_payload is not None + obj._failure_handler_helper_fn(exception=exc, traceback_exception="") + assert obj.model_call_details["standard_logging_object"] is first_payload + + other_exc = _raise_and_catch(_ClientError(status_code=429, message="rate limited")) + obj._failure_handler_helper_fn(exception=other_exc, traceback_exception="") + assert obj.model_call_details["standard_logging_object"] is not first_payload diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index d0169963962..a7fbd069e61 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1,3 +1,4 @@ +import base64 from typing import Any, cast import pytest @@ -11,6 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.litellm_core_utils.prompt_templates.factory import ( THOUGHT_SIGNATURE_SEPARATOR, + _bedrock_converse_messages_pt, ) from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( OPENAI_MAX_TOOL_NAME_LENGTH, @@ -3872,6 +3874,75 @@ def test_tool_result_plain_text_unchanged_by_openai_transform(): assert _image_urls_in_user_messages(result) == [] +TOOL_RESULT_PDF_B64 = base64.b64encode(b"%PDF-1.4 minimal regression fixture").decode() + + +def _base64_pdf_block(): + return { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": TOOL_RESULT_PDF_B64}, + } + + +def test_tool_result_single_document_kept_as_pdf_data_url(): + adapter = LiteLLMAnthropicMessagesAdapter() + translated = adapter.translate_anthropic_messages_to_openai( + messages=[ + _anthropic_tool_use_turn("toolu_01"), + _anthropic_tool_result_turn({"toolu_01": [_base64_pdf_block()]}), + ] + ) + + tool_messages = [m for m in translated if m.get("role") == "tool"] + assert len(tool_messages) == 1 + assert tool_messages[0]["content"] == [ + { + "type": "image_url", + "image_url": {"url": f"data:application/pdf;base64,{TOOL_RESULT_PDF_B64}"}, + } + ] + + +def test_tool_result_text_and_document_reach_bedrock_converse_tool_result(): + """Claude Code >= 2.1.245 sends Read-tool PDF output as a document block inside + tool_result; dropping it left bedrock converse models blind to the PDF content.""" + adapter = LiteLLMAnthropicMessagesAdapter() + translated = adapter.translate_anthropic_messages_to_openai( + messages=[ + AnthropicMessagesUserMessageParam(role="user", content="Read pong.pdf"), + _anthropic_tool_use_turn("toolu_01"), + _anthropic_tool_result_turn( + { + "toolu_01": [ + {"type": "text", "text": "PDF file read: pong.pdf (579 bytes)"}, + _base64_pdf_block(), + ] + } + ), + ] + ) + + converse_messages = _bedrock_converse_messages_pt( + messages=translated, + model="anthropic.claude-haiku-4-5-20251001-v1:0", + llm_provider="bedrock_converse", + ) + + tool_results = [ + block["toolResult"] + for message in converse_messages + for block in message["content"] + if "toolResult" in block + ] + assert len(tool_results) == 1 + documents = [part["document"] for part in tool_results[0]["content"] if "document" in part] + assert len(documents) == 1 + assert documents[0]["format"] == "pdf" + assert documents[0]["source"]["bytes"] == TOOL_RESULT_PDF_B64 + texts = [part["text"] for part in tool_results[0]["content"] if "text" in part] + assert texts == ["PDF file read: pong.pdf (579 bytes)"] + + def test_translate_anthropic_to_openai_carries_prompt_cache_breakpoint_on_system_and_user_blocks(): explicit = {"mode": "explicit"} openai_request, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py index 12ab536ed45..c9170efd18a 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py @@ -42,7 +42,7 @@ def test_reasoning_effort_maps_to_output_config_for_adaptive_model( ) assert "reasoning_effort" not in result - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": expected_effort} diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index f33bb3dda8b..652c1f077a9 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -11,6 +11,10 @@ from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterato BaseAnthropicMessagesStreamingIterator, _incomplete_stream_error_sse_event, _is_message_stop_chunk, + _is_provider_error_chunk, + anthropic_messages_response_as_sse_events, + is_anthropic_content_delta_chunk, + parse_anthropic_error_event, ) @@ -157,6 +161,96 @@ def test_is_message_stop_chunk_ignores_substring_in_payload(): assert _is_message_stop_chunk(delta_frame_with_substring) is False +def test_parse_anthropic_error_event_from_dict_chunk(): + """Regression for #24004: dict-shaped error chunks parse to + (type, message, status) so the Router can decide whether to fall back.""" + chunk = {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}} + assert parse_anthropic_error_event(chunk) == ("overloaded_error", "Overloaded", 503) + assert _is_provider_error_chunk(chunk) is True + + +def test_parse_anthropic_error_event_from_sse_bytes(): + """Regression for #24004: a raw `event: error` SSE frame (what a native + Anthropic/Bedrock passthrough forwards verbatim today) must parse + identically to the dict shape so the Router can raise a fallback.""" + sse_chunk = ( + b"event: error\n" + b'data: {"type": "error", "error": {"type": "internal_server_error", "message": "boom"}}\n\n' + ) + assert parse_anthropic_error_event(sse_chunk) == ("internal_server_error", "boom", 500) + assert _is_provider_error_chunk(sse_chunk) is True + + +def test_parse_anthropic_error_event_defaults_status_for_unknown_type(): + chunk = {"type": "error", "error": {"type": "some_future_error_type", "message": "?"}} + assert parse_anthropic_error_event(chunk) == ("some_future_error_type", "?", 500) + + +def test_parse_anthropic_error_event_missing_message_falls_back_to_type(): + chunk = {"type": "error", "error": {"type": "overloaded_error"}} + assert parse_anthropic_error_event(chunk) == ("overloaded_error", "overloaded_error", 503) + + +def test_parse_anthropic_error_event_non_string_error_type_returns_none(): + """A malformed error body whose `type` field isn't a string (e.g. an + upstream bug sends null or a number) must not be treated as an error + event rather than crashing or forwarding a garbage error_type.""" + chunk = {"type": "error", "error": {"type": None, "message": "boom"}} + assert parse_anthropic_error_event(chunk) is None + + +def test_decoded_sse_data_line_swallows_invalid_json(): + """A `data:` line that isn't valid JSON (a malformed/truncated frame) + must not be treated as an error event or raise, just be ignored.""" + malformed_frame = b"event: error\ndata: {not valid json\n\n" + assert parse_anthropic_error_event(malformed_frame) is None + assert _is_provider_error_chunk(malformed_frame) is False + + +class TestIsAnthropicContentDeltaChunk: + def test_dict_content_block_delta(self): + assert is_anthropic_content_delta_chunk({"type": "content_block_delta"}) is True + + def test_dict_other_type(self): + assert is_anthropic_content_delta_chunk({"type": "message_start"}) is False + + def test_bytes_content_block_delta(self): + assert is_anthropic_content_delta_chunk(b"event: content_block_delta\ndata: {}\n\n") is True + + def test_bytes_other_event(self): + assert is_anthropic_content_delta_chunk(b"event: message_start\ndata: {}\n\n") is False + + def test_neither_dict_nor_bytes(self): + assert is_anthropic_content_delta_chunk("content_block_delta") is False + assert is_anthropic_content_delta_chunk(None) is False + + +@pytest.mark.parametrize( + "chunk", + [ + {"type": "content_block_delta", "delta": {"type": "text_delta", "text": "hi"}}, + b'event: content_block_delta\ndata: {"type": "content_block_delta"}\n\n', + b"raw-bytes", + "error", + None, + ], +) +def test_parse_anthropic_error_event_non_error_chunks_return_none(chunk): + assert parse_anthropic_error_event(chunk) is None + assert _is_provider_error_chunk(chunk) is False + + +def test_parse_anthropic_error_event_ignores_substring_in_payload(): + """A content_block_delta whose partial_json happens to contain the + literal string `"type": "error"` must not be misread as an error event.""" + delta_frame_with_substring = ( + b"event: content_block_delta\n" + b'data: {"type": "content_block_delta", "delta": ' + b'{"type": "input_json_delta", "partial_json": "\\"type\\": \\"error\\""}}\n\n' + ) + assert parse_anthropic_error_event(delta_frame_with_substring) is None + + @pytest.mark.asyncio async def test_async_sse_wrapper_emits_error_when_bytes_stream_only_mentions_message_stop_in_payload(): """ @@ -307,3 +401,117 @@ def test_incomplete_stream_error_sse_event_is_valid_anthropic_error(): "error": {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE}, } assert event.endswith("\n\n") + + +def _decode_sse_events(events: tuple[bytes, ...]) -> list[tuple[str, dict]]: + decoded = [] + for event in events: + assert isinstance(event, bytes) + lines = event.decode().split("\n") + assert lines[0].startswith("event: ") + decoded.append((lines[0].removeprefix("event: "), json.loads(lines[1].removeprefix("data: ")))) + return decoded + + +def test_anthropic_messages_response_as_sse_events_text_block(): + response = { + "id": "msg_1", + "model": "claude-haiku", + "role": "assistant", + "type": "message", + "stop_reason": "end_turn", + "stop_sequence": None, + "content": [{"type": "text", "text": "hello"}], + "usage": {"input_tokens": 3, "output_tokens": 2}, + } + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + types = [event_type for event_type, _ in decoded] + assert types == [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ] + # message_start must not carry generated content itself, matching a real + # streaming response - it arrives via the content_block_delta that follows. + assert decoded[0][1]["message"]["content"] == [] + assert decoded[0][1]["message"]["id"] == "msg_1" + # Bugbot regression: message_start must not carry the completed response's + # final stop_reason/stop_sequence/output_tokens - a real stream keeps those + # null/zero until message_delta, so a client could otherwise treat the + # message as already finished, or double-count output tokens. + assert decoded[0][1]["message"]["stop_reason"] is None + assert decoded[0][1]["message"]["stop_sequence"] is None + assert decoded[0][1]["message"]["usage"] == {"input_tokens": 3, "output_tokens": 0} + assert decoded[1][1]["content_block"] == {"type": "text", "text": ""} + assert decoded[2][1]["delta"] == {"type": "text_delta", "text": "hello"} + assert decoded[4][1]["delta"]["stop_reason"] == "end_turn" + assert decoded[4][1]["usage"] == {"input_tokens": 3, "output_tokens": 2} + + +def test_anthropic_messages_response_as_sse_events_tool_use_block(): + response = { + "id": "msg_2", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "NYC"}}], + "stop_reason": "tool_use", + } + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + content_block_start = dict(decoded)["content_block_start"] + assert content_block_start["content_block"] == { + "type": "tool_use", + "id": "toolu_1", + "name": "get_weather", + "input": {}, + } + content_block_delta = dict(decoded)["content_block_delta"] + assert json.loads(content_block_delta["delta"]["partial_json"]) == {"city": "NYC"} + assert content_block_delta["delta"]["type"] == "input_json_delta" + + +def test_anthropic_messages_response_as_sse_events_thinking_block_emits_signature_delta(): + """Bugbot regression: a thinking block's real `signature` must reach the + client via a trailing signature_delta, not be silently dropped - Anthropic + rejects a replayed assistant message (a follow-up turn, a tool-use + continuation) whose thinking block lacks its original signature.""" + response = { + "id": "msg_5", + "content": [{"type": "thinking", "thinking": "let me think", "signature": "sig-abc123"}], + "stop_reason": "end_turn", + } + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + deltas = [payload["delta"] for event_type, payload in decoded if event_type == "content_block_delta"] + assert deltas == [ + {"type": "thinking_delta", "thinking": "let me think"}, + {"type": "signature_delta", "signature": "sig-abc123"}, + ] + + +def test_anthropic_messages_response_as_sse_events_thinking_block_without_signature_omits_delta(): + response = { + "id": "msg_6", + "content": [{"type": "thinking", "thinking": "let me think", "signature": None}], + "stop_reason": "end_turn", + } + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + deltas = [payload["delta"] for event_type, payload in decoded if event_type == "content_block_delta"] + assert deltas == [{"type": "thinking_delta", "thinking": "let me think"}] + + +def test_anthropic_messages_response_as_sse_events_multiple_blocks_are_indexed(): + response = { + "id": "msg_3", + "content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}], + } + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + starts = [payload for event_type, payload in decoded if event_type == "content_block_start"] + assert [s["index"] for s in starts] == [0, 1] + deltas = [payload for event_type, payload in decoded if event_type == "content_block_delta"] + assert [d["delta"]["text"] for d in deltas] == ["a", "b"] + + +def test_anthropic_messages_response_as_sse_events_no_content_blocks(): + response = {"id": "msg_4", "content": [], "stop_reason": "end_turn"} + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + assert [event_type for event_type, _ in decoded] == ["message_start", "message_delta", "message_stop"] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 8225e7cff39..56f106e407c 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -16,7 +16,10 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, ) -from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT_IMAGE_BOUNDARY +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + TOOL_RESULT_IMAGE_BOUNDARY, + TOOL_RESULT_IMAGE_PLACEHOLDER, +) from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, ) @@ -1555,6 +1558,217 @@ class TestToolResultImages: assert self._input_images(items) == [] +class TestToolResultDocuments: + """Documents inside tool_result blocks must survive translation (LIT-6135): + the function_call_output output becomes a list of parts carrying the joined + text as input_text and each document as an input_file. Without documents the + output stays the plain string it always was.""" + + PDF_B64 = "JVBERi0xLjQKJSBQT05H" + PDF_DATA_URI = "data:application/pdf;base64,JVBERi0xLjQKJSBQT05H" + PDF_URL = "https://example.com/report.pdf" + PNG_B64 = "iVBORw0KGgoAAAANSUhEUg==" + + def _messages(self, tool_result_content): + return [ + {"role": "user", "content": "read the pdf"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_01", "name": "read", "input": {}}], + }, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content} + ], + }, + ] + + def _translate(self, tool_result_content): + return _ADAPTER.translate_messages_to_responses_input(self._messages(tool_result_content)) + + @staticmethod + def _tool_output(items): + return next(item for item in items if item.get("type") == "function_call_output")["output"] + + def _base64_document(self, **extra): + return { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": self.PDF_B64}, + **extra, + } + + def test_text_and_base64_document_produce_part_list(self): + output = self._tool_output( + self._translate([{"type": "text", "text": "PDF file read: mystery.pdf"}, self._base64_document()]) + ) + assert output == [ + {"type": "input_text", "text": "PDF file read: mystery.pdf"}, + {"type": "input_file", "filename": "document.pdf", "file_data": self.PDF_DATA_URI}, + ] + + def test_document_only_produces_single_file_part(self): + output = self._tool_output(self._translate([self._base64_document()])) + assert output == [{"type": "input_file", "filename": "document.pdf", "file_data": self.PDF_DATA_URI}] + + def test_document_title_becomes_filename(self): + output = self._tool_output(self._translate([self._base64_document(title="quarterly-report.pdf")])) + assert output == [ + {"type": "input_file", "filename": "quarterly-report.pdf", "file_data": self.PDF_DATA_URI} + ] + + def test_url_document_becomes_file_url_part(self): + output = self._tool_output( + self._translate([{"type": "document", "source": {"type": "url", "url": self.PDF_URL}}]) + ) + assert output == [{"type": "input_file", "file_url": self.PDF_URL}] + + def test_document_with_empty_data_falls_back_to_string_output(self): + output = self._tool_output( + self._translate( + [ + {"type": "text", "text": "PDF file read"}, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": ""}}, + ] + ) + ) + assert output == "PDF file read" + + def test_document_without_source_dict_keeps_string_output(self): + output = self._tool_output( + self._translate([{"type": "text", "text": "stub"}, {"type": "document", "source": self.PDF_URL}]) + ) + assert output == "stub" + + def test_text_only_tool_result_keeps_plain_string_output(self): + output = self._tool_output(self._translate([{"type": "text", "text": "plain result"}])) + assert output == "plain result" + + def test_file_id_source_document_keeps_string_output(self): + output = self._tool_output( + self._translate( + [ + {"type": "text", "text": "stub"}, + {"type": "document", "source": {"type": "file", "file_id": "file_abc123"}}, + ] + ) + ) + assert output == "stub" + + def test_url_source_without_url_keeps_string_output(self): + output = self._tool_output( + self._translate([{"type": "text", "text": "stub"}, {"type": "document", "source": {"type": "url"}}]) + ) + assert output == "stub" + + def test_text_image_and_document_mix(self): + items = self._translate( + [ + {"type": "text", "text": "captured"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": self.PNG_B64}}, + self._base64_document(), + ] + ) + + output = self._tool_output(items) + assert output == [ + {"type": "input_text", "text": f"captured\n{TOOL_RESULT_IMAGE_PLACEHOLDER}"}, + {"type": "input_file", "filename": "document.pdf", "file_data": self.PDF_DATA_URI}, + ] + + image_message = next( + item + for item in items + if item.get("type") == "message" + and any(part.get("type") == "input_image" for part in item.get("content", [])) + ) + assert image_message["content"] == [ + {"type": "input_text", "text": TOOL_RESULT_IMAGE_BOUNDARY}, + {"type": "input_image", "image_url": f"data:image/png;base64,{self.PNG_B64}"}, + ] + + +class TestUserContentDocuments: + """Documents in plain user content must survive translation (LIT-6144): each + document block becomes an input_file part of the user message, in block order, + exactly like image blocks become input_image parts. Untranslatable documents + are dropped without disturbing the surrounding parts.""" + + PDF_B64 = "JVBERi0xLjQKJSBQT05H" + PDF_DATA_URI = "data:application/pdf;base64,JVBERi0xLjQKJSBQT05H" + PDF_URL = "https://example.com/report.pdf" + EXPLICIT = {"mode": "explicit"} + + def _translate(self, user_content): + return _ADAPTER.translate_messages_to_responses_input([{"role": "user", "content": user_content}]) + + @staticmethod + def _user_content(items): + return next(item for item in items if item.get("type") == "message" and item.get("role") == "user")["content"] + + def _base64_document(self, **extra): + return { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": self.PDF_B64}, + **extra, + } + + def test_document_then_text_keeps_block_order(self): + content = self._user_content( + self._translate([self._base64_document(), {"type": "text", "text": "what does the pdf say?"}]) + ) + assert content == [ + {"type": "input_file", "filename": "document.pdf", "file_data": self.PDF_DATA_URI}, + {"type": "input_text", "text": "what does the pdf say?"}, + ] + + def test_document_title_becomes_filename(self): + content = self._user_content(self._translate([self._base64_document(title="quarterly-report.pdf")])) + assert content == [ + {"type": "input_file", "filename": "quarterly-report.pdf", "file_data": self.PDF_DATA_URI} + ] + + def test_url_document_becomes_file_url_part(self): + content = self._user_content( + self._translate([{"type": "document", "source": {"type": "url", "url": self.PDF_URL}}]) + ) + assert content == [{"type": "input_file", "file_url": self.PDF_URL}] + + def test_document_only_content_still_produces_user_message(self): + content = self._user_content(self._translate([self._base64_document()])) + assert content == [{"type": "input_file", "filename": "document.pdf", "file_data": self.PDF_DATA_URI}] + + def test_empty_base64_data_drops_only_the_document_part(self): + content = self._user_content( + self._translate( + [ + {"type": "text", "text": "still here"}, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": ""}}, + ] + ) + ) + assert content == [{"type": "input_text", "text": "still here"}] + + def test_non_dict_source_drops_only_the_document_part(self): + content = self._user_content( + self._translate([{"type": "text", "text": "still here"}, {"type": "document", "source": self.PDF_URL}]) + ) + assert content == [{"type": "input_text", "text": "still here"}] + + def test_document_breakpoint_rides_on_the_file_part(self): + content = self._user_content( + self._translate([self._base64_document(prompt_cache_breakpoint=self.EXPLICIT)]) + ) + assert content == [ + { + "type": "input_file", + "filename": "document.pdf", + "file_data": self.PDF_DATA_URI, + "prompt_cache_breakpoint": self.EXPLICIT, + } + ] + + def _contains_key(value, key) -> bool: if isinstance(value, dict): return key in value or any(_contains_key(v, key) for v in value.values()) diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py b/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py index ef74249ca8e..288817dff07 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py @@ -5,6 +5,8 @@ Verifies that reasoning_effort=None returns None for all models, including Claude Opus 4.6. """ +import pytest + from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -35,6 +37,16 @@ class TestMapReasoningEffort: ) assert result["type"] == "adaptive" + @pytest.mark.parametrize("effort", ["low", "medium", "high"]) + def test_adaptive_mapping_requests_summarized_display(self, effort): + """Regression LIT-5714: adaptive thinking without ``display`` makes Anthropic + return a blank thinking block, so reasoning_effort callers always got + ``reasoning_content: ""``.""" + result = AnthropicConfig._map_reasoning_effort( + reasoning_effort=effort, model="claude-opus-4-6", custom_llm_provider="anthropic" + ) + assert result["display"] == "summarized" + def test_other_model_low_returns_enabled_with_budget(self): result = AnthropicConfig._map_reasoning_effort( reasoning_effort="low", model="claude-4-sonnet-20250514", custom_llm_provider="anthropic" 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 d4fcbc823a6..11a727c9635 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 @@ -201,6 +201,90 @@ def test_azure_model_router_response_shows_actual_model(): ) +def test_azure_model_router_stamps_selected_model_on_hidden_params(): + """ + The selected model must be stamped on _hidden_params, not left for downstream code to + re-derive by looking for "model-router" in the model string. Deployments whose alias + does not contain that text are invisible to the string check. + """ + from httpx import Response + + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + AzureFoundryModelInfo, + ) + from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj + from litellm.types.utils import ModelResponse + + raw_response_json = { + "id": "chatcmpl-test456", + "object": "chat.completion", + "created": 1234567890, + "model": "grok-4-1-fast-reasoning", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "pong"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + + mock_response = MagicMock(spec=Response) + mock_response.json.return_value = raw_response_json + mock_response.text = json.dumps(raw_response_json) + mock_response.headers = {} + + logging_obj = MagicMock(spec=LiteLLMLoggingObj) + logging_obj.post_call = MagicMock() + logging_obj.model_call_details = {} + + result = AzureModelRouterConfig().transform_response( + model="smart-pick", + raw_response=mock_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={}, + messages=[{"role": "user", "content": "Reply with just pong"}], + optional_params={}, + litellm_params={"model": "azure_ai/model_router/smart-pick"}, + encoding=None, + api_key="test-key", + json_mode=False, + ) + + assert result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] == result.model + assert ( + result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] + == "azure_ai/grok-4-1-fast-reasoning" + ) + assert AzureFoundryModelInfo.get_model_router_selected_model( + result._hidden_params + ) == ("azure_ai/grok-4-1-fast-reasoning") + assert ( + AzureFoundryModelInfo.is_model_router_call( + model="smart-pick", hidden_params=result._hidden_params + ) + is True + ) + + +def test_azure_model_router_stamp_does_not_leak_across_responses(): + """ + ModelResponse declares _hidden_params as a class-level dict, so the stamp has to be written + as a fresh dict. Mutating in place would bleed the selected model into unrelated responses. + """ + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + ) + from litellm.types.utils import ModelResponse + + untouched = ModelResponse() + + assert AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY not in (untouched._hidden_params or {}) + + def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name(): """ Regression test: Azure AI returns 400 when tools contain copilot_mcp_server_name. diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index 53a432427d3..326edde743d 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -341,7 +341,7 @@ def test_messages_thinking_shape_follows_exact_azure_entry_flag(local_model_cost ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem( diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 604f3414775..4d2c077b548 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -366,6 +366,96 @@ def test_output_config_effort_forwarded_into_additional_request_fields(model): assert additional.get("output_config") == {"effort": "high"} +def test_reasoning_effort_requests_summarized_display_converse(): + """Regression LIT-5714: adaptive thinking synthesized from reasoning_effort must + request the summarized display, otherwise the provider returns a blank thinking + block and reasoning_content is always empty.""" + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={"reasoning_effort": "high"}, + optional_params={}, + model="bedrock/converse/us.anthropic.claude-opus-4-7", + drop_params=False, + ) + + assert optional_params["thinking"]["type"] == "adaptive" + assert optional_params["thinking"]["display"] == "summarized" + + +def test_thinking_request_adds_output_tokens_details_response_path(): + """Regression LIT-5714: the Converse usage block has no thinking-token field, so + thinking requests must ask for ``/usage/output_tokens_details`` via + ``additionalModelResponseFieldPaths``.""" + config = AmazonConverseConfig() + + result = config._transform_request( + model="bedrock/converse/us.anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "maxTokens": 256, + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "high"}, + }, + litellm_params={}, + headers={}, + ) + + assert result["additionalModelResponseFieldPaths"] == ("/usage/output_tokens_details",) + + +def test_request_without_thinking_omits_response_field_paths(): + config = AmazonConverseConfig() + + result = config._transform_request( + model="bedrock/converse/us.anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "hi"}], + optional_params={"maxTokens": 256}, + litellm_params={}, + headers={}, + ) + + assert "additionalModelResponseFieldPaths" not in result + + +def test_transform_usage_prefers_provider_reasoning_tokens(): + """Regression LIT-5714: provider-reported thinking tokens must win over the + token_counter estimate derived from visible reasoning text.""" + config = AmazonConverseConfig() + + usage = config.transform_usage( + {"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042}, + reasoning_content="a short reasoning summary", + thinking_ran=True, + provider_reasoning_tokens=1033, + ) + + assert usage.completion_tokens_details.reasoning_tokens == 1033 + assert usage.completion_tokens_details.text_tokens == 3002 - 1033 + + +def test_transform_usage_falls_back_to_estimate_without_provider_tokens(): + config = AmazonConverseConfig() + + usage = config.transform_usage( + {"inputTokens": 40, "outputTokens": 300, "totalTokens": 340}, + reasoning_content="a short reasoning summary", + thinking_ran=True, + ) + + assert usage.completion_tokens_details.reasoning_tokens > 0 + assert usage.completion_tokens_details.reasoning_tokens < 300 + + +def test_thinking_tokens_parsed_from_additional_model_response_fields(): + parsed = AmazonConverseConfig.thinking_tokens_from_additional_fields( + {"usage": {"output_tokens_details": {"thinking_tokens": 92}}} + ) + assert parsed == 92 + assert AmazonConverseConfig.thinking_tokens_from_additional_fields(None) is None + assert AmazonConverseConfig.thinking_tokens_from_additional_fields({"usage": {}}) is None + + @pytest.mark.parametrize( "model,effort,expected_effort", [ diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index cd305d8ed26..4bef59842f1 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -208,6 +208,29 @@ def test_bedrock_converse_streaming_consistent_id(): ), "All chunk IDs must match the one captured from the messageStart event" +def test_converse_streaming_usage_uses_provider_thinking_tokens(): + """Regression LIT-5714: the messageStop event carries provider thinking tokens + under ``additionalModelResponseFields``; the usage chunk must report them instead + of a token_counter estimate.""" + chunks = [ + { + "contentBlockIndex": 0, + "delta": {"reasoningContent": {"text": "thinking about it"}}, + }, + { + "stopReason": "end_turn", + "additionalModelResponseFields": {"usage": {"output_tokens_details": {"thinking_tokens": 1033}}}, + }, + {"usage": {"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042}}, + ] + + decoder = AWSEventStreamDecoder(model="bedrock/anthropic.claude-opus-4-7") + parsed = [decoder.converse_chunk_parser(chunk) for chunk in chunks] + + usage = parsed[-1].usage + assert usage.completion_tokens_details.reasoning_tokens == 1033 + + @pytest.mark.asyncio async def test_make_call_does_not_rechunk_stream_by_default(): """Re-chunking the event stream into fixed 1024-byte blocks holds small diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index d3c28302bf9..1e09afd6919 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -1387,7 +1387,7 @@ def test_bedrock_messages_maps_reasoning_effort_for_adaptive_model( ) assert "reasoning_effort" not in result - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": expected_effort} @@ -2935,7 +2935,7 @@ def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag( ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem(litellm.model_cost[model], "supports_adaptive_thinking", False) diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index f9760cd1dcc..2ea61b5e978 100644 --- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -78,7 +78,7 @@ def test_bedrock_rerank_header_forwarding_sync(model): with ( patch.object(client, "post") as mock_post, - patch( + patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info, ), @@ -171,7 +171,7 @@ async def test_bedrock_rerank_header_forwarding_async(model): with ( patch.object(client, "post", new_callable=AsyncMock) as mock_post, - patch( + patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info, ), @@ -242,7 +242,7 @@ def test_bedrock_rerank_timeout_sync(): with ( patch.object(client, "post") as mock_post, - patch( + patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info, ), @@ -286,7 +286,7 @@ async def test_bedrock_rerank_timeout_async(): with ( patch.object(client, "post", new_callable=AsyncMock) as mock_post, - patch( + patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info, ), @@ -341,7 +341,7 @@ def test_bedrock_rerank_extra_headers_and_headers_merge(): with ( patch.object(client, "post") as mock_post, - patch( + patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info, ), @@ -461,3 +461,32 @@ def test_bedrock_rerank_signs_with_sigv4_even_when_bedrock_api_key_is_set(monkey assert authorization.startswith("AWS4-HMAC-SHA256"), ( f"rerank must sign with SigV4, got Authorization={authorization[:30]}" ) + + +@pytest.mark.asyncio +async def test_bedrock_rerank_records_llm_api_duration(): + """The bedrock rerank handler must feed httpx timing into the logging obj, so the + proxy can emit x-litellm-overhead-duration-ms / x-litellm-timing-* on /rerank.""" + import httpx + + def handle(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=bedrock_rerank_response) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + with patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport + "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", + return_value=create_mock_credentials(), + ): + response = await litellm.arerank( + model="bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0", + query=test_query, + documents=test_documents, + top_n=3, + client=client, + aws_region_name="us-east-1", + ) + + assert response._hidden_params["litellm_overhead_time_ms"] is not None + assert response._hidden_params["_response_ms"] >= response._hidden_params["litellm_overhead_time_ms"] diff --git a/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py b/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py new file mode 100644 index 00000000000..8c6eda605ca --- /dev/null +++ b/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py @@ -0,0 +1,197 @@ +import json +from unittest.mock import MagicMock, patch + +import httpx +import pytest +from botocore.credentials import Credentials + +from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig +from litellm.llms.bedrock_mantle.passthrough.transformation import BedrockMantlePassthroughConfig +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.passthrough.main import llm_passthrough_route +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +MANTLE_API_BASE = "https://bedrock-mantle.us-east-2.api.aws" +INVOKE_ENDPOINT = "model/us.openai.gpt-5.6-sol/invoke" +CONVERSE_ENDPOINT = "model/us.openai.gpt-5.6-sol/converse" +REQUEST_BODY = {"messages": [{"role": "user", "content": "say pong"}], "max_completion_tokens": 64} + + +@pytest.fixture +def no_ambient_aws(monkeypatch): + for name in ( + "AWS_BEARER_TOKEN_BEDROCK", + "BEDROCK_MANTLE_API_KEY", + "BEDROCK_MANTLE_API_BASE", + "BEDROCK_MANTLE_REGION", + "AWS_BEDROCK_RUNTIME_ENDPOINT", + "AWS_REGION_NAME", + "AWS_REGION", + "AWS_DEFAULT_REGION", + ): + monkeypatch.delenv(name, raising=False) + + +def test_bedrock_mantle_registers_its_own_bedrock_passthrough_config(): + config = ProviderConfigManager.get_provider_passthrough_config( + model="us.openai.gpt-5.6-sol", provider=LlmProviders.BEDROCK_MANTLE + ) + assert isinstance(config, BedrockMantlePassthroughConfig) + assert isinstance(config, BedrockPassthroughConfig) + + +def test_mantle_api_base_only_lends_its_region_to_the_runtime_url(no_ambient_aws): + url, base_url = BedrockMantlePassthroughConfig().get_complete_url( + api_base=MANTLE_API_BASE, + api_key=None, + model="us.openai.gpt-5.6-sol", + endpoint=INVOKE_ENDPOINT, + request_query_params=None, + litellm_params={"api_base": MANTLE_API_BASE}, + ) + assert str(url) == f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}" + assert base_url == "https://bedrock-runtime.us-east-2.amazonaws.com" + + +def test_explicit_region_and_non_mantle_api_base_are_kept(no_ambient_aws): + vpc_endpoint = "https://vpce-0123.bedrock-runtime.us-east-1.vpce.amazonaws.com" + url, base_url = BedrockMantlePassthroughConfig().get_complete_url( + api_base=vpc_endpoint, + api_key=None, + model="us.openai.gpt-5.6-sol", + endpoint=INVOKE_ENDPOINT, + request_query_params=None, + litellm_params={"api_base": vpc_endpoint, "aws_region_name": "us-east-1"}, + ) + assert str(url) == f"{vpc_endpoint}/{INVOKE_ENDPOINT}" + assert base_url == vpc_endpoint + + +def test_region_falls_back_to_the_mantle_default_without_any_hint(no_ambient_aws): + url, _ = BedrockMantlePassthroughConfig().get_complete_url( + api_base=None, + api_key=None, + model="us.openai.gpt-5.6-sol", + endpoint=INVOKE_ENDPOINT, + request_query_params=None, + litellm_params={}, + ) + assert str(url) == f"https://bedrock-runtime.us-east-1.amazonaws.com/{INVOKE_ENDPOINT}" + + +@pytest.mark.parametrize( + ("litellm_params", "env", "expected_bearer"), + [ + ({"api_key": "deployment-bedrock-api-key"}, {}, "deployment-bedrock-api-key"), + ({}, {"BEDROCK_MANTLE_API_KEY": "mantle-env-key"}, "mantle-env-key"), + ({}, {"AWS_BEARER_TOKEN_BEDROCK": "aws-env-key"}, "aws-env-key"), + ], +) +def test_sign_request_uses_the_deployment_bearer_token(no_ambient_aws, monkeypatch, litellm_params, env, expected_bearer): + for name, value in env.items(): + monkeypatch.setenv(name, value) + headers, body = BedrockMantlePassthroughConfig().sign_request( + headers={}, + litellm_params=litellm_params, + request_data=REQUEST_BODY, + api_base=f"https://bedrock-runtime.us-east-1.amazonaws.com/{INVOKE_ENDPOINT}", + model="us.openai.gpt-5.6-sol", + ) + assert headers["Authorization"] == f"Bearer {expected_bearer}" + assert body is not None + assert json.loads(body) == REQUEST_BODY + + +def test_sign_request_falls_back_to_sigv4_scoped_to_the_mantle_region(no_ambient_aws): + config = BedrockMantlePassthroughConfig() + with patch.object(config, "get_credentials", return_value=Credentials("AKIA", "secret")): + headers, body = config.sign_request( + headers={}, + litellm_params={"api_base": MANTLE_API_BASE}, + request_data=REQUEST_BODY, + api_base=f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}", + model="us.openai.gpt-5.6-sol", + ) + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIA/") + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + assert body is not None + assert json.loads(body) == REQUEST_BODY + + +@pytest.mark.parametrize( + ("route_kwargs", "env", "expected_bearer"), + [ + ({"api_key": "deployment-bedrock-api-key"}, {}, "deployment-bedrock-api-key"), + ({}, {"BEDROCK_MANTLE_API_KEY": "mantle-env-key"}, "mantle-env-key"), + ], +) +def test_invoke_passthrough_route_reaches_bedrock_runtime_for_a_mantle_deployment( + no_ambient_aws, monkeypatch, route_kwargs, env, expected_bearer +): + for name, value in env.items(): + monkeypatch.setenv(name, value) + client = HTTPHandler() + with ( + patch.object(client.client, "send", return_value=MagicMock(status_code=200)), + patch.object(client.client, "build_request", wraps=client.client.build_request) as build_request, + ): + response = llm_passthrough_route( + model="bedrock_mantle/us.openai.gpt-5.6-sol", + endpoint=INVOKE_ENDPOINT, + method="POST", + api_base=MANTLE_API_BASE, + json=dict(REQUEST_BODY), + client=client, + litellm_logging_obj=MagicMock(), + **route_kwargs, + ) + assert response.status_code == 200 + sent = build_request.call_args.kwargs + assert str(sent["url"]) == f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}" + assert sent["headers"]["Authorization"] == f"Bearer {expected_bearer}" + assert json.loads(sent["content"]) == REQUEST_BODY + + +def _logged_model_response(endpoint, body): + request = httpx.Request("POST", f"https://bedrock-runtime.us-east-1.amazonaws.com/{endpoint}") + return BedrockMantlePassthroughConfig().logging_non_streaming_response( + model="us.openai.gpt-5.6-sol", + custom_llm_provider="bedrock_mantle", + httpx_response=httpx.Response(200, json=body, request=request), + request_data={"messages": [{"role": "user", "content": [{"text": "say pong"}]}]}, + logging_obj=MagicMock(), + endpoint=endpoint, + ) + + +def test_converse_logging_parses_the_converse_response_shape(): + result = _logged_model_response( + CONVERSE_ENDPOINT, + { + "metrics": {"latencyMs": 800.0}, + "output": {"message": {"content": [{"text": "pong"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 8, "outputTokens": 5, "totalTokens": 13}, + }, + ) + assert result.choices[0].message.content == "pong" + assert result.usage.prompt_tokens == 8 + assert result.usage.completion_tokens == 5 + + +def test_invoke_logging_parses_the_openai_chat_response_shape(): + result = _logged_model_response( + INVOKE_ENDPOINT, + { + "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "pong", "role": "assistant"}}], + "created": 1787677792, + "id": "chatcmpl-regression", + "model": "us.openai.gpt-5.6-sol", + "object": "chat.completion", + "usage": {"completion_tokens": 5, "prompt_tokens": 8, "total_tokens": 13}, + }, + ) + assert result.choices[0].message.content == "pong" + assert result.usage.prompt_tokens == 8 + assert result.usage.completion_tokens == 5 diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 9e05d48a18f..00a319f99f0 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -8,7 +8,9 @@ gate, the URL construction for both paths, and the shared Bearer auth. """ import copy - +import json +import logging +from pathlib import Path import pytest from botocore.exceptions import ( @@ -623,6 +625,181 @@ class TestBedrockMantleCodexAdditionalTools: assert "additional_tools" in str(mock_debug.call_args) +class TestBedrockMantleCodexInputItemNormalization: + """Mantle 400s ("Invalid 'input': value did not match any expected variant") + on the Codex history item types agent_message, context_compaction, and + local_shell_call (verified against bedrock-mantle.us-east-1.api.aws with + openai.gpt-5.6-sol), so the config must rewrite them into supported + equivalents. agent_message is what every Codex multi-agent v2 session sends, + and its encrypted_content slot carries the verbatim plaintext payload when + the upstream model never issued encrypted args, so that slot must be + preserved, not dropped. Mantle also rejects assistant messages with + input_text content, so the rewrite must use output_text.""" + + _USER_MESSAGE = { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Continue."}], + } + + def _transform(self, input): + cfg = BedrockMantleResponsesAPIConfig() + return cfg.transform_responses_api_request( + model="openai.gpt-5.6-sol", + input=input, + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + def test_plaintext_agent_message_becomes_assistant_output_text_message(self): + body = self._transform( + input=[ + self._USER_MESSAGE, + { + "type": "agent_message", + "id": "amsg_1", + "author": "/root/arithmetic", + "recipient": "/root", + "content": [{"type": "input_text", "text": "Message Type: FINAL_ANSWER\nPayload:\n2+2 is 4."}], + }, + ] + ) + assert body["input"] == [ + self._USER_MESSAGE, + { + "type": "message", + "role": "assistant", + "content": ({"type": "output_text", "text": "Message Type: FINAL_ANSWER\nPayload:\n2+2 is 4."},), + }, + ] + + def test_agent_message_encrypted_content_payload_is_preserved(self): + body = self._transform( + input=[ + { + "type": "agent_message", + "author": "/root", + "recipient": "/root/arithmetic", + "content": [ + {"type": "input_text", "text": "Message Type: NEW_TASK\nPayload:\n"}, + {"type": "encrypted_content", "encrypted_content": "Answer the question 'what is 2+2'."}, + ], + }, + self._USER_MESSAGE, + ] + ) + assert body["input"][0] == { + "type": "message", + "role": "assistant", + "content": ( + { + "type": "output_text", + "text": "Message Type: NEW_TASK\nPayload:\nAnswer the question 'what is 2+2'.", + }, + ), + } + + def test_agent_message_without_any_text_is_dropped(self): + body = self._transform( + input=[ + {"type": "agent_message", "author": "/root", "recipient": "/root/a", "content": []}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [self._USER_MESSAGE] + + def test_context_compaction_becomes_compaction_with_same_ciphertext(self): + body = self._transform( + input=[ + {"type": "context_compaction", "id": "cc_1", "encrypted_content": "smry_abc123"}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [ + {"type": "compaction", "encrypted_content": "smry_abc123"}, + self._USER_MESSAGE, + ] + + def test_context_compaction_without_ciphertext_is_dropped(self): + body = self._transform( + input=[ + {"type": "context_compaction", "id": "cc_1"}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [self._USER_MESSAGE] + + def test_local_shell_call_becomes_function_call_keeping_call_id_pairing(self): + body = self._transform( + input=[ + { + "type": "local_shell_call", + "id": "lsh_1", + "call_id": "call_1", + "status": "completed", + "action": {"type": "exec", "command": ["echo", "hi"]}, + }, + {"type": "function_call_output", "call_id": "call_1", "output": "hi\n"}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [ + { + "type": "function_call", + "call_id": "call_1", + "name": "local_shell", + "arguments": '{"type": "exec", "command": ["echo", "hi"]}', + }, + {"type": "function_call_output", "call_id": "call_1", "output": "hi\n"}, + self._USER_MESSAGE, + ] + + def test_local_shell_call_without_call_id_is_dropped(self): + body = self._transform( + input=[ + {"type": "local_shell_call", "status": "completed", "action": {"type": "exec", "command": ["ls"]}}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [self._USER_MESSAGE] + + def test_mantle_supported_item_types_pass_through_untouched(self): + supported_items = [ + self._USER_MESSAGE, + {"type": "compaction", "encrypted_content": "smry_abc123"}, + {"type": "function_call", "name": "shell", "arguments": "{}", "call_id": "call_2"}, + {"type": "function_call_output", "call_id": "call_2", "output": "ok"}, + {"type": "tool_search_call", "call_id": "call_3", "execution": "server", "arguments": {"query": "x"}}, + {"type": "tool_search_output", "call_id": "call_3", "status": "completed", "execution": "server", "tools": []}, + {"type": "compaction_trigger"}, + ] + body = self._transform(input=copy.deepcopy(supported_items)) + assert body["input"] == supported_items + + def test_string_input_passes_through(self): + body = self._transform(input="Say hi.") + assert body["input"] == "Say hi." + + def test_rewrite_is_logged_as_warning_naming_the_types(self, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + body = self._transform( + input=[ + {"type": "agent_message", "author": "a", "recipient": "b", "content": [{"type": "input_text", "text": "hi"}]}, + self._USER_MESSAGE, + ] + ) + assert body["input"][0]["role"] == "assistant" + rewrite_warnings = [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "rewrote Codex input item type" in record.getMessage() + ] + assert rewrite_warnings == [ + "Bedrock Mantle Responses API: rewrote Codex input item type(s) ['agent_message'] that Mantle rejects." + ] + + class TestBedrockMantleResponsesRegistry: def test_registry_returns_config_for_gpt_5_5(self, local_cost_map): # gpt-5.x advertises /v1/responses in supported_endpoints (capability) @@ -1523,7 +1700,7 @@ class TestBedrockMantleResponsesPricing: assert info["cache_creation_input_token_cost"] == pytest.approx(cache_creation_cost) assert info["cache_read_input_token_cost"] == pytest.approx(cache_read_cost) assert info["output_cost_per_token"] == pytest.approx(output_cost) - assert info["max_input_tokens"] == 1000000 + assert info["max_input_tokens"] == 1050000 assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2) assert info["cache_creation_input_token_cost_above_272k_tokens"] == pytest.approx(cache_creation_cost * 2) assert info["cache_read_input_token_cost_above_272k_tokens"] == pytest.approx(cache_read_cost * 2) @@ -1565,3 +1742,42 @@ class TestBedrockMantleResponsesPricing: def test_models_registered(self, local_cost_map): assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models + + +def _repo_cost_map(map_name: str) -> dict[str, dict[str, object]]: + repo_root = Path(__file__).resolve().parents[4] + paths = { + "root": repo_root / "model_prices_and_context_window.json", + "bundled_backup": repo_root / "litellm" / "model_prices_and_context_window_backup.json", + } + return json.loads(paths[map_name].read_text()) + + +class TestGpt56MantleRegistryEntries: + """Locks the gpt-5.6 frontier entries to Bedrock Mantle's live behavior. + + Mantle enforces a 1,050,000-token prompt maximum for gpt-5.6 sol/terra/luna + (oversize requests 400 with "prompt tokens (N) exceed model maximum + (1050000)", and a 1,030,590-token request completes), matching the OpenAI + Bedrock guide. mode must stay "responses": Mantle's native + /v1/chat/completions rejects function tools unless reasoning_effort is + "none", so chat traffic has to keep bridging to the Responses API + (see the responses_api_bridge tests above). + """ + + @pytest.mark.parametrize("map_name", ("root", "bundled_backup")) + @pytest.mark.parametrize( + "key", + ( + "bedrock_mantle/openai.gpt-5.6-sol", + "bedrock_mantle/openai.gpt-5.6-terra", + "bedrock_mantle/openai.gpt-5.6-luna", + ), + ) + def test_entry_matches_mantle_enforced_limits(self, map_name, key): + entry = _repo_cost_map(map_name)[key] + assert entry["max_input_tokens"] == 1050000 + assert entry["max_output_tokens"] == 128000 + assert entry["mode"] == "responses" + assert entry["use_openai_responses_path"] is True + assert entry["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"] diff --git a/tests/test_litellm/llms/cerebras/__init__.py b/tests/test_litellm/llms/cerebras/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py b/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py new file mode 100644 index 00000000000..09718b1e6e0 --- /dev/null +++ b/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py @@ -0,0 +1,61 @@ +from litellm.llms.cerebras.chat import CerebrasConfig + + +def test_max_retries_in_supported_params() -> None: + config = CerebrasConfig() + params = config.get_supported_openai_params(model="llama-3.3-70b") + assert "max_retries" in params, ( + f"max_retries must be in CerebrasConfig.get_supported_openai_params(); got: {params!r}" + ) + + +def test_extra_headers_in_supported_params() -> None: + config = CerebrasConfig() + params = config.get_supported_openai_params(model="llama-3.3-70b") + assert "extra_headers" in params, ( + f"extra_headers must be in CerebrasConfig.get_supported_openai_params(); got: {params!r}" + ) + + +def test_core_openai_params_still_supported() -> None: + config = CerebrasConfig() + params = config.get_supported_openai_params(model="llama-3.3-70b") + for expected in ( + "max_tokens", + "max_completion_tokens", + "response_format", + "seed", + "stop", + "stream", + "temperature", + "top_p", + "tool_choice", + "tools", + "user", + ): + assert expected in params, f"{expected!r} unexpectedly missing from Cerebras supported params: {params!r}" + + +def test_map_openai_params_preserves_max_retries() -> None: + config = CerebrasConfig() + result = config.map_openai_params( + non_default_params={"max_retries": 0, "temperature": 0.7}, + optional_params={}, + model="llama-3.3-70b", + drop_params=False, + ) + assert result.get("max_retries") == 0, f"map_openai_params must preserve max_retries=0; got: {result!r}" + assert result.get("temperature") == 0.7 + + +def test_map_openai_params_preserves_max_retries_zero_falsy() -> None: + config = CerebrasConfig() + result = config.map_openai_params( + non_default_params={"max_retries": 0}, + optional_params={}, + model="llama-3.3-70b", + drop_params=False, + ) + assert "max_retries" in result and result["max_retries"] == 0, ( + f"max_retries=0 (falsy) must not be silently omitted; got: {result!r}" + ) diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index a93b14d45f3..694a01cda5f 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2528,6 +2528,37 @@ def test_only_callbacks_that_can_charge_a_frame_are_collected_for_ws_quota(monke assert _collect_ws_project_quota_callbacks() == (quota,) +@pytest.mark.asyncio +async def test_async_rerank_records_llm_api_duration(): + """arerank must feed the httpx timing into the logging obj, so the proxy can emit + x-litellm-overhead-duration-ms / x-litellm-timing-* on /rerank.""" + + def handle(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "rerank-1", + "results": [{"index": 0, "relevance_score": 0.9}], + "meta": {"api_version": {"version": "2"}, "billed_units": {"search_units": 1}}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.arerank( + model="cohere/rerank-v3.5", + query="what is the capital of france", + documents=["paris", "berlin"], + top_n=1, + api_key="fake-key", + client=client, + ) + + assert response._hidden_params["litellm_overhead_time_ms"] is not None + assert response._hidden_params["_response_ms"] >= response._hidden_params["litellm_overhead_time_ms"] + + class _JSONBodyVideoConfig(OpenAIVideoConfig): def use_multipart_form_data(self) -> bool: return False diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index d279b119efe..a35b75a6106 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -1309,3 +1309,79 @@ def test_responses_gpt54_allow_temperature_effort_none( drop_params=False, ) assert params["temperature"] == 0.7 + + +@pytest.mark.parametrize("model", ["gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna"]) +def test_gpt5_6_forwards_reasoning_effort_max_for_the_responses_bridge(config: OpenAIConfig, model: str): + """A chat request carrying tools or a reasoning summary is converted to /v1/responses further + down main.py, and that surface accepts max. This runs before litellm has decided to bridge, so + refusing max here would break the cursor thinking-max shape that works today. Plain chat + completions still answer max with a provider 400, and the capability list below is what keeps + the level out of the picker.""" + params = config.map_openai_params( + non_default_params={"reasoning_effort": "max"}, + optional_params={}, + model=model, + drop_params=False, + ) + assert params["reasoning_effort"] == "max" + + +@pytest.mark.parametrize("model", ["gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna"]) +def test_gpt5_6_never_advertises_reasoning_effort_max(model: str): + """/v1/chat/completions answers max with "Unsupported value: 'reasoning_effort' does not support + 'max' with this model. Supported values are: 'none', 'low', 'medium', 'high', and 'xhigh'", so no + gpt-5.6 entry asserts supports_max_reasoning_effort and the advertised set stops at xhigh.""" + from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts + + resolved = resolve_supported_reasoning_efforts(litellm.get_model_info(model), deployment_is_mapped=True) + assert resolved is not None + assert "max" not in resolved + assert "xhigh" in resolved + + +def test_gpt5_6_keeps_reasoning_effort_max_on_the_responses_api( + responses_config: OpenAIResponsesAPIConfig, +): + """/v1/responses accepts max for gpt-5.6, and that is the surface the cursor thinking-max + variant resolves onto, so the responses path keeps carrying the level chat refuses.""" + params = responses_config.map_openai_params( + response_api_optional_params=ResponsesAPIOptionalRequestParams( + reasoning={"effort": "max"}, + ), + model="gpt-5.6", + drop_params=False, + ) + assert params["reasoning"] == {"effort": "max"} + + +def test_gpt5_forwards_levels_the_chat_gate_does_not_own(config: OpenAIConfig): + """Only xhigh is gated on this surface. max reaches the provider (or the responses bridge) and + is answered there, which is what happened before per-group capabilities existed.""" + params = config.map_openai_params( + non_default_params={"reasoning_effort": "max"}, + optional_params={}, + model="gpt-5.1", + drop_params=False, + ) + assert params["reasoning_effort"] == "max" + + +def test_gpt5_rejects_xhigh_for_models_without_the_flag(config: OpenAIConfig): + with pytest.raises(litellm.utils.UnsupportedParamsError): + config.map_openai_params( + non_default_params={"reasoning_effort": "xhigh"}, + optional_params={}, + model="gpt-5.1", + drop_params=False, + ) + + +def test_gpt5_drops_xhigh_when_requested(config: OpenAIConfig): + params = config.map_openai_params( + non_default_params={"reasoning_effort": "xhigh"}, + optional_params={}, + model="gpt-5.1", + drop_params=True, + ) + assert "reasoning_effort" not in params diff --git a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py new file mode 100644 index 00000000000..0b9fd5364f9 --- /dev/null +++ b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py @@ -0,0 +1,387 @@ +import json +import logging +from unittest.mock import MagicMock + +import httpx +import pytest + +import litellm +from litellm.exceptions import UnsupportedParamsError +from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj +from litellm.llms.openai.chat.gpt_transformation import ( + OpenAIChatCompletionStreamingHandler, +) +from litellm.llms.together_ai.chat.transformation import TogetherAIChatConfig +from litellm.types.utils import LlmProviders, ModelResponse + +TOOL_CALLING_MODEL = "openai/gpt-oss-20b" +REASONING_MODEL = "deepseek-ai/DeepSeek-V3.1" +UNMAPPED_MODEL = "example-org/brand-new-model" +NO_TOOLS_MODEL = "example-org/no-tools-model" + +TOOL_PARAMS = ("tools", "tool_choice", "function_call") + +WEATHER_TOOLS = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}] + + +@pytest.fixture(autouse=True) +def force_local_model_cost(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map + + monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url)) + + +@pytest.fixture +def registry_disables_function_calling(monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + f"together_ai/{NO_TOOLS_MODEL}", + {"litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": False}, + ) + + +@pytest.fixture +def together_warning_log(caplog): + from litellm._logging import verbose_logger + + verbose_logger.addHandler(caplog.handler) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + yield caplog + verbose_logger.removeHandler(caplog.handler) + + +def test_supported_params_tool_calling_model(): + supported = TogetherAIChatConfig().get_supported_openai_params(model=TOOL_CALLING_MODEL) + + for param in (*TOOL_PARAMS, "response_format"): + assert param in supported + + +def test_supported_params_unmapped_model_keeps_tool_params(): + supported = TogetherAIChatConfig().get_supported_openai_params(model=UNMAPPED_MODEL) + + for param in TOOL_PARAMS: + assert param in supported + assert "response_format" not in supported + assert "stream" in supported + assert "temperature" in supported + + +def test_supported_params_no_tools_model_keeps_tool_params(registry_disables_function_calling): + supported = TogetherAIChatConfig().get_supported_openai_params(model=NO_TOOLS_MODEL) + + for param in TOOL_PARAMS: + assert param in supported + assert "response_format" not in supported + + +def test_map_openai_params_tool_calling_model_passes_tools(): + mapped = TogetherAIChatConfig().map_openai_params( + non_default_params={"tools": WEATHER_TOOLS, "tool_choice": "auto"}, + optional_params={}, + model=TOOL_CALLING_MODEL, + drop_params=False, + ) + + assert mapped["tools"] == WEATHER_TOOLS + assert mapped["tool_choice"] == "auto" + + +@pytest.mark.parametrize("drop_params", [False, True]) +def test_map_openai_params_unmapped_model_passes_tools_through(drop_params, together_warning_log): + mapped = TogetherAIChatConfig().map_openai_params( + non_default_params={"tools": WEATHER_TOOLS, "tool_choice": "required"}, + optional_params={}, + model=UNMAPPED_MODEL, + drop_params=drop_params, + ) + + assert mapped["tools"] == WEATHER_TOOLS + assert mapped["tool_choice"] == "required" + assert UNMAPPED_MODEL in together_warning_log.text + assert "passing tools, tool_choice through" in together_warning_log.text + + +def test_map_openai_params_no_tools_model_drops_tools_with_warning( + registry_disables_function_calling, together_warning_log +): + mapped = TogetherAIChatConfig().map_openai_params( + non_default_params={"tools": WEATHER_TOOLS, "temperature": 0.5}, + optional_params={}, + model=NO_TOOLS_MODEL, + drop_params=True, + ) + + assert "tools" not in mapped + assert mapped["temperature"] == 0.5 + assert NO_TOOLS_MODEL in together_warning_log.text + assert "dropping tools" in together_warning_log.text + + +def test_map_openai_params_no_tools_model_raises_without_drop_params(registry_disables_function_calling): + with pytest.raises(UnsupportedParamsError, match="does not support parameters"): + TogetherAIChatConfig().map_openai_params( + non_default_params={"tools": WEATHER_TOOLS}, + optional_params={}, + model=NO_TOOLS_MODEL, + drop_params=False, + ) + + +def test_map_openai_params_reasoning_model_passes_sampling_params(): + mapped = TogetherAIChatConfig().map_openai_params( + non_default_params={"temperature": 0.2, "max_tokens": 512}, + optional_params={}, + model=REASONING_MODEL, + drop_params=False, + ) + + assert mapped["temperature"] == 0.2 + assert mapped["max_tokens"] == 512 + + +def test_map_openai_params_drops_text_response_format(): + mapped = TogetherAIChatConfig().map_openai_params( + non_default_params={"response_format": {"type": "text"}, "temperature": 0.5}, + optional_params={}, + model=REASONING_MODEL, + drop_params=False, + ) + + assert "response_format" not in mapped + assert mapped["temperature"] == 0.5 + + +def test_map_openai_params_keeps_json_response_format(): + response_format = {"type": "json_object"} + + mapped = TogetherAIChatConfig().map_openai_params( + non_default_params={"response_format": response_format}, + optional_params={}, + model=TOOL_CALLING_MODEL, + drop_params=False, + ) + + assert mapped["response_format"] == response_format + + +def _transform_response(message: dict) -> ModelResponse: + raw_response_json = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1234567890, + "model": REASONING_MODEL, + "choices": [{"index": 0, "message": message, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = raw_response_json + mock_response.text = json.dumps(raw_response_json) + mock_response.headers = {} + logging_obj = MagicMock(spec=LiteLLMLoggingObj) + logging_obj.post_call = MagicMock() + logging_obj.model_call_details = {} + + return TogetherAIChatConfig().transform_response( + model=REASONING_MODEL, + raw_response=mock_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={}, + messages=[{"role": "user", "content": "What is 2+2?"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key="test-key", + json_mode=False, + ) + + +def test_transform_response_maps_reasoning_to_reasoning_content(): + result = _transform_response( + {"role": "assistant", "content": "4", "reasoning": "2+2 equals 4"} + ) + + assert result.choices[0].message.content == "4" + assert result.choices[0].message.reasoning_content == "2+2 equals 4" + + +def test_transform_response_preserves_reasoning_content_field(): + result = _transform_response( + {"role": "assistant", "content": "4", "reasoning_content": "adding 2 and 2"} + ) + + assert result.choices[0].message.reasoning_content == "adding 2 and 2" + + +def test_streaming_chunk_maps_delta_reasoning_to_reasoning_content(): + iterator = TogetherAIChatConfig().get_model_response_iterator( + streaming_response=iter(()), sync_stream=True + ) + assert isinstance(iterator, OpenAIChatCompletionStreamingHandler) + + parsed = iterator.chunk_parser( + { + "id": "chunk-1", + "created": 1234567890, + "model": REASONING_MODEL, + "choices": [{"index": 0, "delta": {"reasoning": "thinking about 2+2"}}], + } + ) + + assert parsed.choices[0]["delta"]["reasoning_content"] == "thinking about 2+2" + + +def test_streaming_chunk_preserves_tool_call_index_and_id(): + iterator = TogetherAIChatConfig().get_model_response_iterator( + streaming_response=iter(()), sync_stream=True + ) + + def parse_tool_call_chunk(tool_call: dict): + parsed = iterator.chunk_parser( + { + "id": "chunk-1", + "created": 1234567890, + "model": TOOL_CALLING_MODEL, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "", "tool_calls": [tool_call]}}], + } + ) + return parsed.choices[0]["delta"]["tool_calls"][0] + + opener = parse_tool_call_chunk( + { + "index": 1, + "id": "call_abc123", + "type": "function", + "function": {"name": "get_weather", "arguments": ""}, + } + ) + continuation = parse_tool_call_chunk( + {"index": 1, "id": "", "type": "function", "function": {"arguments": '{"city": "San'}} + ) + + assert opener["index"] == 1 + assert opener["id"] == "call_abc123" + assert opener["function"]["name"] == "get_weather" + assert continuation["index"] == 1 + assert continuation["function"]["arguments"] == '{"city": "San' + + +def test_together_ai_config_alias_points_at_chat_config(): + assert litellm.TogetherAIConfig is litellm.TogetherAIChatConfig + config = litellm.TogetherAIConfig(max_tokens=10) + assert isinstance(config, TogetherAIChatConfig) + + +def test_provider_config_manager_returns_together_chat_config(): + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_chat_config( + model=REASONING_MODEL, provider=LlmProviders.TOGETHER_AI + ) + + assert isinstance(config, TogetherAIChatConfig) + + +def test_completion_routes_through_together_chat_config(): + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + captured_requests = [] + + def respond(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-together", + "object": "chat.completion", + "created": 1234567890, + "model": REASONING_MODEL, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "4", + "reasoning": "2+2 equals 4", + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + ) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + + response = litellm.completion( + model=f"together_ai/{REASONING_MODEL}", + messages=[{"role": "user", "content": "What is 2+2?"}], + api_key="fake-key", + client=client, + ) + + request = captured_requests[0] + assert str(request.url) == "https://api.together.ai/v1/chat/completions" + assert request.headers["authorization"] == "Bearer fake-key" + assert json.loads(request.content)["model"] == REASONING_MODEL + assert response.choices[0].message.content == "4" + assert response.choices[0].message.reasoning_content == "2+2 equals 4" + + +def test_completion_unmapped_model_sends_tools_to_together(): + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + captured_requests = [] + + def respond(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-together-tools", + "object": "chat.completion", + "created": 1234567890, + "model": UNMAPPED_MODEL, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "San Francisco"}', + }, + } + ], + }, + "finish_reason": "tool_calls", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + ) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + + response = litellm.completion( + model=f"together_ai/{UNMAPPED_MODEL}", + messages=[{"role": "user", "content": "What is the weather in San Francisco?"}], + tools=WEATHER_TOOLS, + tool_choice="auto", + api_key="fake-key", + client=client, + ) + + request_body = json.loads(captured_requests[0].content) + assert request_body["tools"] == WEATHER_TOOLS + assert request_body["tool_choice"] == "auto" + tool_call = response.choices[0].message.tool_calls[0] + assert tool_call.function.name == "get_weather" + assert json.loads(tool_call.function.arguments) == {"city": "San Francisco"} diff --git a/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py b/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py new file mode 100644 index 00000000000..3364b1b3872 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py @@ -0,0 +1,231 @@ +import pytest + +import litellm +from litellm.interactions.utils import get_provider_interactions_api_config +from litellm.llms.gemini.interactions.transformation import ( + GoogleAIStudioInteractionsConfig, +) +from litellm.llms.vertex_ai.interactions.transformation import ( + VertexAIInteractionsConfig, +) +from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + +GLOBAL_BASE = "https://aiplatform.googleapis.com/v1beta1/projects/test-proj/locations/global/interactions" + + +class MinterRecorder: + def __init__(self, resolved_project: str = "creds-proj") -> None: + self.calls: list[tuple[VERTEX_CREDENTIALS_TYPES | None, str | None]] = [] + self.resolved_project = resolved_project + + def __call__( + self, + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + ) -> tuple[str, str]: + self.calls.append((credentials, project_id)) + return "test-token", project_id or self.resolved_project + + +@pytest.fixture +def minter(): + return MinterRecorder() + + +@pytest.fixture +def config(minter): + return VertexAIInteractionsConfig(mint_access_token=minter) + + +@pytest.fixture +def litellm_params(): + return GenericLiteLLMParams(vertex_project="test-proj", vertex_credentials="creds.json") + + +class TestRegistration: + def test_vertex_ai_returns_vertex_config(self): + assert isinstance(get_provider_interactions_api_config("vertex_ai"), VertexAIInteractionsConfig) + + def test_vertex_ai_beta_returns_vertex_config(self): + assert isinstance(get_provider_interactions_api_config("vertex_ai_beta"), VertexAIInteractionsConfig) + + def test_gemini_still_returns_google_ai_studio_config(self): + gemini_config = get_provider_interactions_api_config("gemini") + assert isinstance(gemini_config, GoogleAIStudioInteractionsConfig) + assert not isinstance(gemini_config, VertexAIInteractionsConfig) + + def test_lazy_import_resolves(self): + assert litellm.VertexAIInteractionsConfig is VertexAIInteractionsConfig + + def test_custom_llm_provider_is_vertex_ai(self, config): + assert config.custom_llm_provider == LlmProviders.VERTEX_AI + + +class TestValidateEnvironment: + def test_sets_bearer_auth_without_gemini_headers(self, config, minter, litellm_params): + headers = config.validate_environment( + headers={}, + model="gemini-omni-flash-preview", + litellm_params=litellm_params, + ) + + assert headers["Authorization"] == "Bearer test-token" + assert headers["Content-Type"] == "application/json" + assert "x-goog-api-key" not in headers + assert "Api-Revision" not in headers + assert minter.calls == [("creds.json", "test-proj")] + + def test_caller_authorization_wins(self, config, litellm_params): + headers = config.validate_environment( + headers={"Authorization": "Bearer caller-token"}, + model="gemini-omni-flash-preview", + litellm_params=litellm_params, + ) + + assert headers["Authorization"] == "Bearer caller-token" + + +class TestGetCompleteUrl: + def test_defaults_to_global_v1beta1(self, config, litellm_params): + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params=dict(litellm_params), + ) + + assert url == GLOBAL_BASE + + def test_stream_appends_alt_sse(self, config, litellm_params): + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params=dict(litellm_params), + stream=True, + ) + + assert url == f"{GLOBAL_BASE}?alt=sse" + + def test_multi_region_location_uses_rep_host(self, config): + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_project": "test-proj", "vertex_location": "us"}, + ) + + assert url == "https://aiplatform.us.rep.googleapis.com/v1beta1/projects/test-proj/locations/us/interactions" + + def test_regional_location_uses_regional_host(self, config): + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_project": "test-proj", "vertex_location": "us-central1"}, + ) + + assert url == ( + "https://us-central1-aiplatform.googleapis.com" + "/v1beta1/projects/test-proj/locations/us-central1/interactions" + ) + + def test_location_env_fallback_is_ignored(self, config, monkeypatch): + monkeypatch.setenv("VERTEXAI_LOCATION", "us-east5") + + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_project": "test-proj"}, + ) + + assert url == GLOBAL_BASE + + def test_api_base_override(self, config, litellm_params): + url = config.get_complete_url( + api_base="https://proxy.example.test", + model="gemini-omni-flash-preview", + litellm_params=dict(litellm_params), + ) + + assert url == "https://proxy.example.test/v1beta1/projects/test-proj/locations/global/interactions" + + def test_project_resolved_from_credentials_when_not_passed(self, config, monkeypatch): + monkeypatch.delenv("VERTEXAI_PROJECT", raising=False) + monkeypatch.setattr(litellm, "vertex_project", None) + + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_credentials": "creds.json"}, + ) + + assert url == "https://aiplatform.googleapis.com/v1beta1/projects/creds-proj/locations/global/interactions" + + def test_invalid_location_rejected(self, config): + with pytest.raises(ValueError, match="Invalid vertex_location"): + config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_project": "test-proj", "vertex_location": "evil.com#"}, + ) + + def test_missing_project_rejected(self, monkeypatch): + monkeypatch.delenv("VERTEXAI_PROJECT", raising=False) + monkeypatch.setattr(litellm, "vertex_project", None) + + def unresolved_minter( + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + ) -> tuple[str, str]: + return "test-token", "" + + with pytest.raises(ValueError, match="Vertex AI project is required"): + VertexAIInteractionsConfig(mint_access_token=unresolved_minter).get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={}, + ) + + +class TestInteractionByIdRequests: + def test_get_url(self, config, litellm_params): + url, request_body = config.transform_get_interaction_request( + interaction_id="abc123", + api_base="", + litellm_params=litellm_params, + headers={}, + ) + + assert url == f"{GLOBAL_BASE}/abc123" + assert request_body == {} + + def test_get_url_encodes_interaction_id(self, config, litellm_params): + url, _ = config.transform_get_interaction_request( + interaction_id="id/with space", + api_base="", + litellm_params=litellm_params, + headers={}, + ) + + assert url == f"{GLOBAL_BASE}/id%2Fwith%20space" + + def test_delete_url(self, config, litellm_params): + url, request_body = config.transform_delete_interaction_request( + interaction_id="abc123", + api_base="", + litellm_params=litellm_params, + headers={}, + ) + + assert url == f"{GLOBAL_BASE}/abc123" + assert request_body == {} + + def test_cancel_url(self, config, litellm_params): + url, request_body = config.transform_cancel_interaction_request( + interaction_id="abc123", + api_base="", + litellm_params=litellm_params, + headers={}, + ) + + assert url == f"{GLOBAL_BASE}/abc123:cancel" + assert request_body == {} diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index ba2f20e2337..f19e169dc9e 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -538,7 +538,7 @@ def test_messages_thinking_shape_follows_exact_vertex_entry_flag(local_model_cos ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem( diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 04f38b5e2ed..7d2e3f45c04 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -47,6 +47,7 @@ from litellm.proxy.auth.auth_checks import ( _virtual_key_soft_budget_check, get_key_object, get_user_object, + invalidate_team_member_spend_state, vector_store_access_check, ) from litellm.caching.in_memory_cache import InMemoryCache @@ -4868,10 +4869,9 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): 3. When team_alias is None, NO alias-key operation happens (no delete of an empty-keyed entry, no spurious write). 4. DELETES the team_id-keyed entry from the internal usage cache - BEFORE the fresh write (LIT-4391). `_get_team_object_from_cache` - consults the internal usage cache first, so a leftover copy there - (backfilled from a Redis shared with `user_api_key_cache`) would - keep serving the pre-update team allowlist. + BEFORE the fresh write (LIT-4391). `_get_team_object_from_cache` no + longer reads the internal usage cache (LIT-5944), but the delete + protects mixed-version rolling deploys where older workers still do. """ from unittest.mock import AsyncMock, MagicMock @@ -4980,8 +4980,10 @@ async def test_team_update_not_shadowed_by_internal_usage_cache_lit_4391(): Regression test for LIT-4391: keys with models=["all-team-models"] kept getting 403 team_model_access_denied for models added via /team/update. - `_get_team_object_from_cache` consults the internal usage cache BEFORE - `user_api_key_cache`. When both share one Redis (enable_redis_auth_cache), + `_get_team_object_from_cache` used to consult the internal usage cache + BEFORE `user_api_key_cache` (removed in LIT-5944; this test now also + guards against reintroducing that read). + When both share one Redis (enable_redis_auth_cache), any team read backfills the internal cache's in-memory tier with the team object. `_cache_team_object` (the /team/update refresh) only wrote `user_api_key_cache`, so that backfilled copy kept shadowing the update @@ -5053,6 +5055,75 @@ async def test_team_update_not_shadowed_by_internal_usage_cache_lit_4391(): ) +class _CountingFakeRedis(_SharedFakeRedis): + """Counts per-key Redis round-trips so tests can pin the number of + network operations a code path issues.""" + + def __init__(self): + super().__init__() + self.get_calls: int = 0 + + async def async_get_cache(self, key, **kwargs): + self.get_calls += 1 + return await super().async_get_cache(key, **kwargs) + + +@pytest.mark.asyncio +async def test_warm_team_object_reads_issue_no_redis_ops_lit_5944(): + """ + Regression test for LIT-5944: project/team-scoped virtual-key requests + paid ~4 awaited Redis GETs per request just to re-read the team object. + + `_get_team_object_from_cache` used to consult + `proxy_logging_obj.internal_usage_cache.dual_cache` (in-memory TTL 1s, + Redis-backed) BEFORE `user_api_key_cache`. Nothing writes team objects + into that internal cache — `_cache_team_object` only DELETES the key + there — so when `user_api_key_cache` has no Redis tier the shared Redis + key stays absent forever and every team lookup in the auth hot path + (4 call sites per chat-completion request) became a guaranteed-miss + Redis round-trip, saturating the event loop at high TPS. + + Pins: once `_cache_team_object` has cached a team, repeated + `get_team_object` reads are served from `user_api_key_cache`'s in-memory + tier and issue ZERO Redis operations. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + from litellm.proxy.auth.auth_checks import _cache_team_object, get_team_object + + team_id = "team-lit-5944" + counting_redis = _CountingFakeRedis() + user_api_key_cache = UserApiKeyCache() + proxy_logging_obj = MagicMock() + proxy_logging_obj.internal_usage_cache.dual_cache = DualCache( + redis_cache=counting_redis, + default_in_memory_ttl=1, + ) + prisma_client = MagicMock() + + await _cache_team_object( + team_id=team_id, + team_table=LiteLLM_TeamTableCachedObj(team_id=team_id, models=["model-a"]), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + for _ in range(4): + team_obj = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + assert team_obj is not None and team_obj.models == ["model-a"] + + assert counting_redis.get_calls == 0, ( + "Warm team-object reads must be served from user_api_key_cache's " + "in-memory tier without any Redis round-trips. " + f"Got {counting_redis.get_calls} Redis GETs for 4 get_team_object calls." + ) + + @pytest.mark.asyncio async def test_cache_team_object_tolerates_cache_invalidation_failures(): """ @@ -6939,3 +7010,387 @@ def test_model_has_no_cost_mapping_alias_to_a_group_priced_through_model_info_is router = _router_with_a_group_priced_through_model_info() assert model_has_no_cost_mapping(model="model-info-priced-alias", llm_router=router) is False + + +@pytest.mark.parametrize( + "user_route, expected", + [ + ("/internal-models/v1/chat/completions", True), + ("/internal-models/newly-registered-model/predict", True), + ("/internal-models-other/v1/chat/completions", False), + ("/anthropic/v1/messages", False), + ], +) +def test_team_allowed_routes_wildcard_prefix_matches_unregistered_passthrough_routes(user_route, expected): + """A `/prefix/*` entry in `team_allowed_routes` must cover every route under that prefix, so + passthrough endpoints registered after the proxy config was written are reachable without an + exact-route config change.""" + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.TEAM, + user_route=user_route, + litellm_proxy_roles=LiteLLM_JWTAuth(team_allowed_routes=["/internal-models/*"]), + ) + is expected + ) + + +def test_team_allowed_routes_exact_route_does_not_become_a_prefix_grant(): + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + roles = LiteLLM_JWTAuth(team_allowed_routes=["/internal-models/model-a"]) + + assert ( + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-a", litellm_proxy_roles=roles) + is True + ) + assert ( + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-b", litellm_proxy_roles=roles) + is False + ) + + +def test_admin_allowed_routes_wildcard_prefix_is_honored(): + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + roles = LiteLLM_JWTAuth(admin_allowed_routes=["/internal-models/*"]) + + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.PROXY_ADMIN, user_route="/internal-models/anything", litellm_proxy_roles=roles + ) + is True + ) + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.PROXY_ADMIN, user_route="/other/anything", litellm_proxy_roles=roles + ) + is False + ) + + +def test_team_allowed_routes_named_route_group_still_resolves(): + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + roles = LiteLLM_JWTAuth(team_allowed_routes=["openai_routes"]) + + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.TEAM, user_route="/v1/chat/completions", litellm_proxy_roles=roles + ) + is True + ) + assert ( + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/key/generate", litellm_proxy_roles=roles) + is False + ) + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_sets_the_spend_counter_and_clears_both_membership_cache_keys(): + """A team-member budget reset (new_spend passed) must SET the spend counter to the reset + value, clear its DB-floor marker, AND invalidate both independently-keyed membership caches + (user_api_key_auth.py's admission check writes one key format, budget_reservation.py and + auth_checks.py's own get_team_membership() write the other) or a stale read keeps 429ing + after the reset. Asserted against real cache reads, not mock call args, so a change that + keeps the call but drops its effect still fails.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_user-1", value="stale-membership") + await real_cache.async_set_cache(key="team_membership:user-1:team-1", value="stale-membership") + + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:user-1:team-1", value=999.0) + real_spend_counter_cache.in_memory_cache.set_cache( + key="spend_db_floor:spend:team_member:user-1:team-1", value=999.0 + ) + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + new_spend=0.0, + ) + + assert await real_cache.async_get_cache(key="team-1_user-1") is None + assert await real_cache.async_get_cache(key="team_membership:user-1:team-1") is None + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 0.0 + assert ( + real_spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:team_member:user-1:team-1") + == 0.0 + ), "the DB-floor marker kept the pre-reset value; a stale-floor read can raise the counter right back up" + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_leaves_the_live_spend_counter_alone_without_new_spend(): + """team_member_update only changes the budget cap, not the tracked spend, so it calls + invalidate_team_member_spend_state with no new_spend. Deleting the live spend counter in that + case would force the next read to reseed from the DB's own spend column, which lags the live + counter via periodic batch writes, briefly UNDER-enforcing the raised cap against a spend + value lower than what was actually tracked (regression: PR #37971 Bugbot finding).""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_user-1", value="stale-membership") + + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:user-1:team-1", value=999.0) + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + ) + + assert await real_cache.async_get_cache(key="team-1_user-1") is None + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 999.0 + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_deleting(): + """/key/{key}/reset_spend SETs its counter to the reset value rather than deleting it, so a + worker's next read reflects it directly instead of falling back through a DB reseed. A reset + caller passing new_spend must match that precedent, not merely delete the counter.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + real_spend_counter_cache = DualCache() + fake_redis_cache = MagicMock() + fake_redis_cache.async_set_cache = AsyncMock() + real_spend_counter_cache.redis_cache = fake_redis_cache + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + new_spend=2.5, + ) + + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 2.5 + fake_redis_cache.async_set_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1", value=2.5, ttl=60) + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_deletes_redis_counter_when_set_fails(): # test-quality-ok: only observable effect is the fallback call on the same fake client + """Redis reads take priority over the local in-memory copy (get_current_spend reads Redis + first), so a failed Redis SET would otherwise leave the OLD pre-reset value authoritative + for every worker even though the reset reported success. On a failed SET, the stale Redis + entry must be deleted instead, so the next read clean-misses and reseeds from the DB.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + real_spend_counter_cache = DualCache() + fake_redis_cache = MagicMock() + fake_redis_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down")) + fake_redis_cache.async_delete_cache = AsyncMock() + real_spend_counter_cache.redis_cache = fake_redis_cache + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + new_spend=2.5, + ) + + fake_redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1") + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_raises_503_when_both_redis_writes_fail(): + """If the Redis SET fails AND the fallback DELETE fails, the stale pre-reset counter is still + authoritative in Redis for every worker. Reporting success would silently keep 429ing the + member, so the reset must surface a 503 instead (regression: PR #37971 Greptile finding).""" + from fastapi import HTTPException + + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + real_spend_counter_cache = DualCache() + fake_redis_cache = MagicMock() + fake_redis_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down")) + fake_redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis still down")) + real_spend_counter_cache.redis_cache = fake_redis_cache + + with ( + patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ), + pytest.raises(HTTPException) as exc_info, + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + new_spend=2.5, + ) + + assert exc_info.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_to_remote_workers(): + """The test above only proves the handling worker's own spend counter is + cleared. A remote worker's spend counter is a separate DualCache instance; + if the reset never reaches it, that worker keeps enforcing the pre-reset + spend the moment its own Redis read for the counter fails and it falls + back to its own (now-stale) in-memory copy. Drives the actual message + published onto the invalidation channel through a second, independent + AuthCacheInvalidationSubscriber standing in for that remote worker, rather + than asserting on the publish call args.""" + from redis.asyncio import Redis + + from litellm.caching.dual_cache import DualCache + from litellm.caching.in_memory_cache import InMemoryCache + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + published: list[tuple[str, str]] = [] + + class _RecordingRedisClient(Redis): + def __init__(self) -> None: + pass + + async def publish(self, channel: str, message: str) -> int: + published.append((channel, message)) + return 1 + + class _FakeRedisCache: + def __init__(self) -> None: + self.namespace = None + + def init_async_client(self) -> object: + return _RecordingRedisClient() + + local_spend_counter_cache = DualCache() + + remote_user_api_key_cache = UserApiKeyCache() + remote_spend_counter_in_memory_cache = InMemoryCache() + remote_spend_counter_in_memory_cache.set_cache("spend:team_member:user-1:team-1", 999.0) + remote_spend_counter_in_memory_cache.set_cache("spend_db_floor:spend:team_member:user-1:team-1", 999.0) + + with ( + patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", local_spend_counter_cache + ), + patch( # test-quality-ok: injects a fake pub/sub-capable redis cache; no live redis in this unit test + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(), + ), + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=UserApiKeyCache(), + new_spend=0.0, + ) + + def _published_message_for(cache_key: str) -> str: + matches = [message for _, message in published if json.loads(message)["cache_key"] == cache_key] + assert matches, f"{cache_key} never reached the cross-worker invalidation channel" + return matches[-1] + + remote_subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(), + user_api_key_cache=remote_user_api_key_cache, + additional_in_memory_caches=(remote_spend_counter_in_memory_cache,), + ) + for cache_key in ("spend:team_member:user-1:team-1", "spend_db_floor:spend:team_member:user-1:team-1"): + remote_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API + {"type": "message", "data": _published_message_for(cache_key)} + ) + + assert remote_spend_counter_in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0 + assert ( + remote_spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0 + ), "the DB-floor marker was not broadcast; a remote worker can re-raise the counter off its stale floor" + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_self_delivered_broadcast_does_not_erase_the_reset(): + """The handling worker subscribes to the same invalidation channel it publishes on, so it + receives its own reset message. A delete-style broadcast would erase the post-reset counter + and floor marker the handler just wrote, reopening the stale-floor race the reset closed + (regression: PR #37971 Greptile finding). The broadcast carries the reset value as a SET, so + applying the self-delivered message must leave both keys at the post-reset value.""" + from redis.asyncio import Redis + + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + published: list[tuple[str, str]] = [] + + class _RecordingRedisClient(Redis): + def __init__(self) -> None: + pass + + async def publish(self, channel: str, message: str) -> int: + published.append((channel, message)) + return 1 + + class _FakeRedisCache: + def __init__(self) -> None: + self.namespace = None + + def init_async_client(self) -> object: + return _RecordingRedisClient() + + local_spend_counter_cache = DualCache() + local_user_api_key_cache = UserApiKeyCache() + + with ( + patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", local_spend_counter_cache + ), + patch( # test-quality-ok: injects a fake pub/sub-capable redis cache; no live redis in this unit test + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(), + ), + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=local_user_api_key_cache, + new_spend=0.0, + ) + + own_subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(), + user_api_key_cache=local_user_api_key_cache, + additional_in_memory_caches=(local_spend_counter_cache.in_memory_cache,), + ) + for _, message in published: + own_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API + {"type": "message", "data": message} + ) + + assert local_spend_counter_cache.in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0, ( + "the handler's self-delivered broadcast erased the post-reset spend counter" + ) + assert ( + local_spend_counter_cache.in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0 + ), "the handler's self-delivered broadcast erased the post-reset floor marker, reopening the stale-floor race" diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index b0094b81112..90b3b29d919 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -701,3 +701,58 @@ async def test_auth_failure_ip_stamp_does_not_mutate_callers_request_data(): ) assert request_data == {"model": "gpt-4o"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth_error,expect_traceback", + [ + pytest.param( + ProxyException( + message="Authentication Error", type=ProxyErrorTypes.auth_error, param=None, code=401 + ), + False, + id="expected_401_no_traceback", + ), + pytest.param(ValueError("unexpected internal error"), True, id="unexpected_error_keeps_traceback"), + ], +) +async def test_handle_authentication_error_traceback_only_for_unexpected_errors(auth_error, expect_traceback, caplog): + """Regression for LIT-6043: expected 4xx auth rejections must not format a + traceback via logger.exception; unexpected errors must keep it.""" + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + verbose_proxy_logger.propagate = True + try: + try: + raise auth_error + except (ProxyException, ValueError) as caught: + with caplog.at_level("ERROR", logger="LiteLLM Proxy"), pytest.raises(ProxyException): + await handler._handle_authentication_error( + caught, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-bad-key", + ) + finally: + verbose_proxy_logger.propagate = False + + records = [r for r in caplog.records if "user_api_key_auth(): Exception occured" in r.getMessage()] + assert len(records) == 1 + assert (records[0].exc_info is not None) is expect_traceback diff --git a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 468e8aabae8..7d5fc1a3544 100644 --- a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -6,6 +6,7 @@ from unittest.mock import patch import pytest from redis.asyncio import Redis +from litellm.caching.in_memory_cache import InMemoryCache from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AUTH_CACHE_INVALIDATION_CHANNEL, AuthCacheInvalidationSubscriber, @@ -144,6 +145,37 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None: assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL] +@pytest.mark.asyncio +async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() -> None: + """ + The spend-counter half of the same cross-worker gap: a remote worker's own + spend counter can hold a stale value (its fallback path when that worker's + own Redis read for the counter fails), and only clearing user_api_key_cache + on message would leave that separate DualCache's in-memory copy untouched. + """ + cache = UserApiKeyCache() + spend_counter_in_memory_cache = InMemoryCache() + spend_counter_in_memory_cache.set_cache("spend:team_member:u-1:t-1", 999.0) + assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is not None + + pubsub = _QueuePubSub(initial_messages=[_invalidation_message("spend:team_member:u-1:t-1")]) + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), + user_api_key_cache=cache, + additional_in_memory_caches=(spend_counter_in_memory_cache,), + ) + subscriber.start() + try: + for _ in range(200): + if spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None: + break + await asyncio.sleep(0.01) + finally: + await subscriber.stop() + + assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None + + @pytest.mark.asyncio async def test_subscriber_ignores_malformed_messages() -> None: cache = UserApiKeyCache() diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py index 9853ce7e1cf..135175dd29d 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py @@ -2,7 +2,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest - from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, Member from litellm.proxy.management_endpoints.scim.scim_transformations import ( ScimTransformations, @@ -95,25 +94,17 @@ def mock_prisma_client(): class TestScimTransformations: @pytest.mark.asyncio - async def test_transform_litellm_user_to_scim_user( - self, mock_user, mock_prisma_client - ): + async def test_transform_litellm_user_to_scim_user(self, mock_user, mock_prisma_client): mock_client, mock_find_unique = mock_prisma_client # Mock the team lookup - team1 = LiteLLM_TeamTable( - team_id="team-1", team_alias="Team One", members_with_roles=[] - ) - team2 = LiteLLM_TeamTable( - team_id="team-2", team_alias="Team Two", members_with_roles=[] - ) + team1 = LiteLLM_TeamTable(team_id="team-1", team_alias="Team One", members_with_roles=[]) + team2 = LiteLLM_TeamTable(team_id="team-2", team_alias="Team Two", members_with_roles=[]) mock_find_unique.side_effect = [team1, team2] with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - mock_user - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(mock_user) assert scim_user.id == mock_user.user_id assert scim_user.userName == mock_user.user_email @@ -129,21 +120,15 @@ class TestScimTransformations: assert scim_user.groups[1].display == "Team Two" @pytest.mark.asyncio - async def test_transform_user_with_scim_metadata( - self, mock_user_with_scim_metadata, mock_prisma_client - ): + async def test_transform_user_with_scim_metadata(self, mock_user_with_scim_metadata, mock_prisma_client): mock_client, mock_find_unique = mock_prisma_client # Mock the team lookup - team1 = LiteLLM_TeamTable( - team_id="team-1", team_alias="Team One", members_with_roles=[] - ) + team1 = LiteLLM_TeamTable(team_id="team-1", team_alias="Team One", members_with_roles=[]) mock_find_unique.return_value = team1 with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - mock_user_with_scim_metadata - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(mock_user_with_scim_metadata) assert scim_user.name.givenName == "Test" assert scim_user.name.familyName == "User" @@ -160,15 +145,11 @@ class TestScimTransformations: teams=[], created_at=None, updated_at=None, - metadata={ - "scim_enterprise": {"costCenter": "CC-42", "department": "Platform"} - }, + metadata={"scim_enterprise": {"costCenter": "CC-42", "department": "Platform"}}, ) with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - user - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user) assert scim_user.enterprise_user is not None assert scim_user.enterprise_user.costCenter == "CC-42" @@ -176,9 +157,7 @@ class TestScimTransformations: assert SCIM_ENTERPRISE_USER_SCHEMA in scim_user.schemas @pytest.mark.asyncio - async def test_transform_user_with_entitlements_and_roles_metadata( - self, mock_prisma_client - ): + async def test_transform_user_with_entitlements_and_roles_metadata(self, mock_prisma_client): mock_client, mock_find_unique = mock_prisma_client mock_find_unique.return_value = None @@ -190,17 +169,13 @@ class TestScimTransformations: created_at=None, updated_at=None, metadata={ - "scim_entitlements": [ - {"value": "jira-software", "display": "Jira Software"} - ], + "scim_entitlements": [{"value": "jira-software", "display": "Jira Software"}], "scim_roles": [{"value": "engineering-admin", "primary": True}], }, ) with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - user - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user) assert scim_user.entitlements is not None assert scim_user.entitlements[0].value == "jira-software" @@ -210,9 +185,7 @@ class TestScimTransformations: assert scim_user.roles[0].primary is True @pytest.mark.asyncio - async def test_transform_user_with_malformed_directory_metadata_fails_soft( - self, mock_prisma_client - ): + async def test_transform_user_with_malformed_directory_metadata_fails_soft(self, mock_prisma_client): """Metadata is writable outside the SCIM surface; a corrupted value on one user must omit the attribute, not fail the whole directory response""" mock_client, mock_find_unique = mock_prisma_client @@ -233,9 +206,7 @@ class TestScimTransformations: ) with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - user - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user) assert scim_user.id == "user-corrupt" assert scim_user.entitlements is None @@ -244,22 +215,14 @@ class TestScimTransformations: assert SCIM_ENTERPRISE_USER_SCHEMA not in scim_user.schemas @pytest.mark.asyncio - async def test_transform_user_without_enterprise_metadata_omits_schema( - self, mock_user, mock_prisma_client - ): + async def test_transform_user_without_enterprise_metadata_omits_schema(self, mock_user, mock_prisma_client): mock_client, mock_find_unique = mock_prisma_client - team1 = LiteLLM_TeamTable( - team_id="team-1", team_alias="Team One", members_with_roles=[] - ) - team2 = LiteLLM_TeamTable( - team_id="team-2", team_alias="Team Two", members_with_roles=[] - ) + team1 = LiteLLM_TeamTable(team_id="team-1", team_alias="Team One", members_with_roles=[]) + team2 = LiteLLM_TeamTable(team_id="team-2", team_alias="Team Two", members_with_roles=[]) mock_find_unique.side_effect = [team1, team2] with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - mock_user - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(mock_user) assert scim_user.enterprise_user is None assert SCIM_ENTERPRISE_USER_SCHEMA not in scim_user.schemas @@ -309,36 +272,28 @@ class TestScimTransformations: assert dumped_attrs["roles"][0]["value"] == "engineering-admin" @pytest.mark.asyncio - async def test_transform_litellm_team_to_scim_group( - self, mock_team, mock_prisma_client - ): + async def test_transform_litellm_team_to_scim_group(self, mock_team, mock_prisma_client): mock_client, _ = mock_prisma_client with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_group = await ScimTransformations.transform_litellm_team_to_scim_group( - mock_team - ) + scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(mock_team) assert scim_group.id == mock_team.team_id assert scim_group.displayName == mock_team.team_alias assert len(scim_group.members) == 2 - assert scim_group.members[0].value == "test@example.com" + assert scim_group.members[0].value == "user-123" assert scim_group.members[0].display == "test@example.com" - assert scim_group.members[1].value == "test2@example.com" + assert scim_group.members[1].value == "user-456" assert scim_group.members[1].display == "test2@example.com" @pytest.mark.asyncio - async def test_transform_team_marks_members_as_users( - self, mock_team, mock_prisma_client - ): + async def test_transform_team_marks_members_as_users(self, mock_team, mock_prisma_client): """A LiteLLM team only holds users, and stating the member type keeps the response from emitting a null ``type`` now that SCIMMember carries one.""" mock_client, _ = mock_prisma_client with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_group = await ScimTransformations.transform_litellm_team_to_scim_group( - mock_team - ) + scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(mock_team) assert [member.type for member in scim_group.members] == ["User", "User"] @@ -351,9 +306,7 @@ class TestScimTransformations: result = ScimTransformations._get_scim_user_name(mock_user_minimal) assert result == ScimTransformations.DEFAULT_SCIM_DISPLAY_NAME - def test_get_scim_family_name( - self, mock_user, mock_user_with_scim_metadata, mock_user_minimal - ): + def test_get_scim_family_name(self, mock_user, mock_user_with_scim_metadata, mock_user_minimal): # User with alias result = ScimTransformations._get_scim_family_name(mock_user) assert result == mock_user.user_alias @@ -366,9 +319,7 @@ class TestScimTransformations: result = ScimTransformations._get_scim_family_name(mock_user_minimal) assert result == ScimTransformations.DEFAULT_SCIM_FAMILY_NAME - def test_get_scim_given_name( - self, mock_user, mock_user_with_scim_metadata, mock_user_minimal - ): + def test_get_scim_given_name(self, mock_user, mock_user_with_scim_metadata, mock_user_minimal): # User with alias result = ScimTransformations._get_scim_given_name(mock_user) assert result == mock_user.user_alias @@ -382,14 +333,10 @@ class TestScimTransformations: assert result == ScimTransformations.DEFAULT_SCIM_NAME def test_get_scim_member_value(self): - # Member with email - member_with_email = Member( - user_id="user-123", user_email="test@example.com", role="admin" - ) + member_with_email = Member(user_id="user-123", user_email="test@example.com", role="admin") result = ScimTransformations._get_scim_member_value(member_with_email) - assert result == member_with_email.user_email + assert result == member_with_email.user_id - # Member without email should fall back to user_id member_without_email = Member(user_id="user-456", user_email=None, role="user") result = ScimTransformations._get_scim_member_value(member_without_email) assert result == member_without_email.user_id @@ -415,9 +362,7 @@ class TestScimTransformations: mock_find_unique.return_value = None with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - user_with_uuid_email - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user_with_uuid_email) assert scim_user.id == user_with_uuid_email.user_id assert scim_user.emails is None or len(scim_user.emails) == 0 @@ -443,9 +388,7 @@ class TestScimTransformations: mock_find_unique.return_value = None with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - user_with_none_email - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user_with_none_email) assert scim_user.id == user_with_none_email.user_id assert scim_user.emails is None or len(scim_user.emails) == 0 diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 5f6c1a2375b..d51e3abde83 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -752,6 +752,63 @@ async def test_handle_existing_user_by_email_syncs_roster_and_dedups_teams(mocke assert update_calls[0].kwargs["data"]["teams"] == ["team-a", "team-b"] +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_without_teams_preserves_memberships(mocker): + """Adoption via POST /Users without ``groups`` must keep the user's existing teams. + + Regression: Entra manages membership exclusively through /Groups and never sends + ``groups`` on POST /Users, so the empty team list was treated as the desired + state and the adopted user was removed from every team roster and had ``teams`` + overwritten with []. + """ + existing_user = mocker.MagicMock() + existing_user.user_id = "adopted-id" + existing_user.user_email = "member@example.com" + existing_user.user_alias = "Member" + existing_user.teams = ["team-a", "team-b"] + existing_user.metadata = {} + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={}) + + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=None), + ) + mock_team_member_add = mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + AsyncMock(), + ) + mock_team_member_delete = mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", + AsyncMock(), + ) + + new_user_request = NewUserRequest( + user_id="entra-object-id", + user_email="member@example.com", + user_alias="Member", + teams=[], + metadata={}, + auto_create_key=False, + ) + + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, new_user_request=new_user_request + ) + + mock_team_member_add.assert_not_awaited() + mock_team_member_delete.assert_not_awaited() + + update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list + assert len(update_calls) == 1 + assert update_calls[0].kwargs["where"] == {"user_id": "adopted-id"} + assert update_calls[0].kwargs["data"]["teams"] == ["team-a", "team-b"] + + @pytest.mark.asyncio async def test_handle_existing_user_by_email_roster_add_failure_blocks_teams_write(mocker): """A genuine roster add failure must propagate and must not persist the teams array. @@ -872,11 +929,16 @@ async def test_handle_existing_user_by_email_roster_remove_failure_blocks_teams_ AsyncMock(side_effect=HTTPException(status_code=500, detail={"error": "No db connected"})), ) + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + AsyncMock(), + ) + new_user_request = NewUserRequest( user_id="uid", user_email="member@example.com", user_alias="Member", - teams=[], + teams=["replacement-team"], metadata={}, auto_create_key=False, ) @@ -917,11 +979,16 @@ async def test_handle_existing_user_by_email_roster_remove_already_absent_is_noo AsyncMock(return_value=None), ) + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + AsyncMock(), + ) + new_user_request = NewUserRequest( user_id="uid", user_email="member@example.com", user_alias="Member", - teams=[], + teams=["replacement-team"], metadata={}, auto_create_key=False, ) @@ -933,7 +1000,7 @@ async def test_handle_existing_user_by_email_roster_remove_already_absent_is_noo mock_team_member_delete.assert_awaited_once() update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list assert len(update_calls) == 1 - assert update_calls[0].kwargs["data"]["teams"] == [] + assert update_calls[0].kwargs["data"]["teams"] == ["replacement-team"] @pytest.mark.asyncio 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 f6d74a189bc..f33854eeb81 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -4,16 +4,18 @@ from contextlib import asynccontextmanager from datetime import datetime, timezone from types import SimpleNamespace from typing import Optional, cast -from unittest.mock import AsyncMock, MagicMock, call, patch +from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch import pytest from fastapi import HTTPException from fastapi.testclient import TestClient +from pydantic import ValidationError from litellm._uuid import uuid from litellm.proxy._types import UserAPIKeyAuth # Import UserAPIKeyAuth from litellm.proxy._types import ( + LiteLLM_BudgetTable, LiteLLM_BudgetTableFull, LiteLLM_ModelTable, LiteLLM_OrganizationMembershipTable, @@ -27,7 +29,9 @@ from litellm.proxy._types import ( Member, ProxyErrorTypes, ProxyException, + ResetSpendRequest, TeamMemberAddRequest, + TeamMemberUpdateRequest, UpdateTeamRequest, ) from litellm.proxy.management_endpoints.team_endpoints import ( @@ -42,15 +46,21 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _transform_teams_to_deleted_records, _update_model_table, _validate_and_populate_member_user_info, + _validate_team_member_reset_spend_value, _verify_team_access, delete_team, list_available_teams, + reset_team_member_spend_fn, router, team_member_add_duplication_check, team_member_delete, + team_member_update, update_team, validate_team_org_change, ) +from litellm.proxy.management_helpers.access_group_team_sync import ( + TEAM_ADVISORY_LOCK_SQL, +) from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, ) @@ -68,7 +78,11 @@ client = TestClient(app) def _wire_team_create_tx(prisma_client): """`/team/new` inserts the team and mirrors it onto the access groups in one transaction, - so a mocked client has to hand its team table back out of `db.tx()`.""" + so a mocked client has to hand its team table back out of `db.tx()`. + + A `/team/new` carrying members then adds them under the team's advisory lock, and those + writes run on that lock's transaction, so `tx()` has to hand back the mocked tables too + for the per-table assertions on `prisma_client.db.*` to keep seeing them.""" @asynccontextmanager async def _tx(): @@ -78,18 +92,67 @@ def _wire_team_create_tx(prisma_client): ) prisma_client.db.tx = lambda *_args, **_kwargs: _tx() + _wire_member_add_tx(prisma_client) + + +def _wire_member_add_tx(prisma_client): + """/team/member_add takes the team's advisory lock, re-reads the roster under it, and runs + the user, budget, and membership writes on that same transaction, so a mocked client has + to hand its own table mocks back out of `tx()`. + + Tables resolve on access, not here, since tests routinely replace `db.` after + wiring the transaction.""" + + class _Tx: + query_raw = AsyncMock(return_value=[{"members_with_roles": []}]) + + def __getattr__(self, table_name): + return getattr(prisma_client.db, table_name) + + tx = _Tx() + tx_cm = MagicMock() + tx_cm.__aenter__ = AsyncMock(return_value=tx) + tx_cm.__aexit__ = AsyncMock(return_value=None) + prisma_client.tx = MagicMock(return_value=tx_cm) def _wire_member_delete_tx(prisma_client): - """/team/member_delete's four cleanups run inside one transaction, so a mocked - client has to hand back its own table mocks out of `tx()` for the existing - per-table assertions to keep seeing the calls.""" + """/team/member_delete's four cleanups, plus the advisory-lock re-read that now guards + them, run inside one transaction, so a mocked client has to hand back its own table + mocks (and a `query_raw` that answers the locked re-read from the same team row the + test already configured on `find_unique`) out of `tx()` for the existing per-table + assertions to keep seeing the calls.""" + + async def _query_raw(sql, team_id): + if sql != TEAM_ADVISORY_LOCK_SQL: + team_row = await prisma_client.db.litellm_teamtable.find_unique(where={"team_id": team_id}) + if team_row is not None: + return [{"members_with_roles": team_row.model_dump()["members_with_roles"]}] + return [] + + class _Tx: + query_raw = staticmethod(_query_raw) + + def __getattr__(self, table_name): + return getattr(prisma_client.db, table_name) + + tx = _Tx() + tx_cm = MagicMock() + tx_cm.__aenter__ = AsyncMock(return_value=tx) + tx_cm.__aexit__ = AsyncMock(return_value=None) + prisma_client.tx = MagicMock(return_value=tx_cm) + + +def _wire_team_delete_tx(prisma_client): + """`/team/delete` deletes the team rows and runs its post-delete reference sweep under + every team's advisory lock in one transaction, so a mocked client has to hand its own + table mocks (and db-level execute_raw) back out of `tx()` for existing per-table + assertions on `prisma_client.db.*` to keep seeing those calls.""" tx = SimpleNamespace( litellm_teamtable=prisma_client.db.litellm_teamtable, - litellm_usertable=prisma_client.db.litellm_usertable, litellm_teammembership=prisma_client.db.litellm_teammembership, - litellm_verificationtoken=prisma_client.db.litellm_verificationtoken, - litellm_deletedverificationtoken=prisma_client.db.litellm_deletedverificationtoken, + query_raw=AsyncMock(return_value=[]), + execute_raw=prisma_client.db.execute_raw, ) tx_cm = MagicMock() tx_cm.__aenter__ = AsyncMock(return_value=tx) @@ -1662,6 +1725,7 @@ async def test_process_team_members_single_member(): default_team_budget_id="budget-123", allowed_models=None, budget_duration=None, + tx=None, ) @@ -1802,8 +1866,8 @@ async def test_update_team_members_list_duplicate_prevention(): async def test_add_team_members_reconciles_against_freshly_locked_row(): """ Regression: _add_team_members_to_team must build the new members_with_roles - from the row it re-reads under a lock inside the write transaction, not from - the stale complete_team_data snapshot captured at the start of the request. + from the row it re-reads under the team's advisory lock, not from the stale + complete_team_data snapshot captured at the start of the request. Two concurrent /team/member_add calls for the same team read the same snapshot; without the locked re-read the losing write rewrites the whole @@ -1864,24 +1928,89 @@ async def test_add_team_members_reconciles_against_freshly_locked_row(): written_ids = sorted(m["user_id"] for m in json.loads(captured["data"]["members_with_roles"])) assert written_ids == ["alice", "bob", "zed"] - lock_reads = [call for call in tx.query_raw.call_args_list if "FOR UPDATE" in str(call.args[0])] - assert lock_reads, "expected a SELECT ... FOR UPDATE row-lock read before the write" + assert tx.query_raw.call_args_list[0].args == (TEAM_ADVISORY_LOCK_SQL, "test-team-lock"), ( + "expected the team's advisory lock to be acquired before the members_with_roles read" + ) + assert not any("FOR UPDATE" in str(call.args[0]) for call in tx.query_raw.call_args_list), ( + "a row lock here can deadlock with the access-group endpoints; only the advisory lock is safe" + ) assert [m.user_id for m in updated_team.members_with_roles] == ["zed", "alice", "bob"] @pytest.mark.asyncio -async def test_add_team_members_cleans_up_when_the_team_is_deleted_mid_request(): +async def test_add_team_members_runs_member_writes_on_the_lock_holding_transaction(): + """ + Regression pin against exhausting the connection pool with advisory-lock waiters. + + Every concurrent /team/member_add for one team holds a pooled connection while it waits + on the team's advisory lock. If the holder's member writes went to the regular client, + it would need a second connection to finish, so enough concurrent adds fill the pool + with waiters and the holder can never commit or release the lock. The member writes + therefore have to run on the transaction that already owns the connection. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + _add_team_members_to_team, + ) + + added_user = MagicMock() + added_user.user_id = "bob" + added_user.model_dump.return_value = {"user_id": "bob", "teams": ["team-pool"]} + created_budget = MagicMock() + created_budget.budget_id = "budget-pool" + membership = MagicMock() + membership.model_dump.return_value = { + "team_id": "team-pool", + "user_id": "bob", + "budget_id": "budget-pool", + "litellm_budget_table": None, + } + + tx = MagicMock() + tx.query_raw = AsyncMock(return_value=[{"members_with_roles": []}]) + tx.litellm_teamtable.update = AsyncMock( + return_value=LiteLLM_TeamTable(team_id="team-pool", members_with_roles=[]) + ) + tx.litellm_usertable.upsert = AsyncMock(return_value=added_user) + tx.litellm_usertable.update_many = AsyncMock() + tx.litellm_budgettable.create = AsyncMock(return_value=created_budget) + tx.litellm_teammembership.create = AsyncMock(return_value=membership) + + tx_cm = MagicMock() + tx_cm.__aenter__ = AsyncMock(return_value=tx) + tx_cm.__aexit__ = AsyncMock(return_value=None) + + prisma_client = MagicMock() + prisma_client.tx = MagicMock(return_value=tx_cm) + type(prisma_client).db = PropertyMock( + side_effect=AssertionError("member writes must not reach for a second pooled connection") + ) + + _, updated_users, updated_team_memberships = await _add_team_members_to_team( + data=TeamMemberAddRequest( + team_id="team-pool", + member=Member(user_id="bob", role="user"), + max_budget_in_team=50.0, + ), + complete_team_data=LiteLLM_TeamTable(team_id="team-pool", members_with_roles=[]), + prisma_client=cast(object, prisma_client), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + litellm_proxy_admin_name="admin", + ) + + assert [user.user_id for user in updated_users] == ["bob"] + assert [tm.budget_id for tm in updated_team_memberships] == ["budget-pool"] + + +@pytest.mark.asyncio +async def test_add_team_members_writes_nothing_when_the_team_is_deleted_mid_request(): """ Regression pin for the /team/member_add vs /team/delete race. - The user row and membership writes land before the reconcile takes the team - row lock, so a /team/delete that commits in between has already run its own - reference sweep and cannot see them. The empty locked SELECT is the only - signal that happened, and leaving it at that would strand the member on a - deleted team id, which authorization paths that trust `user.teams` would - treat as membership if the id were ever recreated. So the request must sweep - the references it just wrote and fail, not report success. + The advisory lock is acquired, and the team is gone, before any write is attempted: + the empty locked SELECT is proof a /team/delete already committed under the same + lock, so this request must fail without writing the user or membership rows in the + first place, rather than writing them and then trying to sweep them back out. """ from litellm.proxy.management_endpoints.team_endpoints import ( _add_team_members_to_team, @@ -1900,9 +2029,10 @@ async def test_add_team_members_cleans_up_when_the_team_is_deleted_mid_request() prisma_client.db.execute_raw = AsyncMock() prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + process_team_members = AsyncMock(return_value=([], [])) with patch( "litellm.proxy.management_endpoints.team_endpoints._process_team_members", - new=AsyncMock(return_value=([], [])), + new=process_team_members, ): with pytest.raises(HTTPException) as exc_info: await _add_team_members_to_team( @@ -1917,14 +2047,10 @@ async def test_add_team_members_cleans_up_when_the_team_is_deleted_mid_request() ) assert exc_info.value.status_code == 404 + process_team_members.assert_not_awaited() tx.litellm_teamtable.update.assert_not_awaited() - - assert prisma_client.db.execute_raw.await_args_list == [ - call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-deleted-mid-add") - ] - prisma_client.db.litellm_teammembership.delete_many.assert_awaited_once_with( - where={"team_id": {"in": ("team-deleted-mid-add",)}} - ) + prisma_client.db.execute_raw.assert_not_awaited() + prisma_client.db.litellm_teammembership.delete_many.assert_not_awaited() def test_add_new_models_to_team_with_existing_models(): @@ -4239,6 +4365,86 @@ async def test_team_member_delete_cleans_verification_tokens( ) +@pytest.mark.asyncio +async def test_team_member_delete_reads_on_the_lock_holding_transaction( + mock_db_client, mock_admin_auth +): + """ + Regression pin against exhausting the connection pool with advisory-lock waiters. + + Every concurrent removal for one team holds a pooled connection while it waits on the + team's advisory lock, and /team/delete fans its per-member removals out concurrently. + A holder whose reads went to the regular client would need a second connection to + finish, so enough waiters fill the pool and the holder can never release the lock. + Both reads therefore have to run on the transaction that already owns the connection. + """ + from litellm.proxy._types import TeamMemberDeleteRequest + from litellm.proxy.management_endpoints.team_endpoints import team_member_delete + + test_team_id = "team-del-pool-123" + test_user_id = "user-del-pool-123" + roster_entry = {"user_id": test_user_id, "user_email": None, "role": "user"} + + mock_team_row = MagicMock() + mock_team_row.model_dump.return_value = { + "team_id": test_team_id, + "members_with_roles": [roster_entry], + "team_member_permissions": [], + "metadata": {}, + "models": [], + "spend": 0.0, + } + mock_db_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_team_row + ) + + user_row = MagicMock() + user_row.user_id = test_user_id + user_row.teams = [test_team_id] + + # Both are wired to answer, so the endpoint completes either way and the awaits below + # are what tells which connection it read on. + pooled_user_read = AsyncMock(return_value=[user_row]) + pooled_token_read = AsyncMock(return_value=[]) + mock_db_client.db.litellm_usertable.find_many = pooled_user_read + mock_db_client.db.litellm_verificationtoken.find_many = pooled_token_read + + tx = MagicMock() + tx.query_raw = AsyncMock(return_value=[{"members_with_roles": [roster_entry]}]) + tx.litellm_teamtable.update = AsyncMock(return_value=mock_team_row) + tx.litellm_usertable.find_many = AsyncMock(return_value=[user_row]) + tx.litellm_usertable.update = AsyncMock() + tx.litellm_teammembership.delete_many = AsyncMock() + tx.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + tx.litellm_verificationtoken.delete_many = AsyncMock() + + tx_cm = MagicMock() + tx_cm.__aenter__ = AsyncMock(return_value=tx) + tx_cm.__aexit__ = AsyncMock(return_value=None) + mock_db_client.tx = MagicMock(return_value=tx_cm) + + await team_member_delete( + data=TeamMemberDeleteRequest(team_id=test_team_id, user_id=test_user_id), + user_api_key_dict=mock_admin_auth, + ) + + tx.litellm_usertable.find_many.assert_awaited_once_with( + where={"user_id": {"in": [test_user_id]}} + ) + tx.litellm_verificationtoken.find_many.assert_awaited_once_with( + where={"user_id": {"in": [test_user_id]}, "team_id": test_team_id} + ) + pooled_user_read.assert_not_awaited() + pooled_token_read.assert_not_awaited() + + tx.litellm_usertable.update.assert_awaited_once_with( + where={"user_id": test_user_id}, data={"teams": {"set": []}} + ) + tx.litellm_teammembership.delete_many.assert_awaited_once_with( + where={"team_id": test_team_id, "user_id": test_user_id} + ) + + @pytest.mark.parametrize( "roster_email", ["Alice@Example.com", "alice-invited-as@example.com"], @@ -7404,6 +7610,7 @@ async def test_delete_team_persists_deleted_teams( mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) mock_tx_cm.__aexit__ = AsyncMock(return_value=False) mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + _wire_team_delete_tx(mock_prisma_client) monkeypatch.setattr( "litellm.proxy.proxy_server.prisma_client", @@ -7474,15 +7681,14 @@ async def test_delete_team_sweeps_references_outside_members_with_roles( cache_state_when_rows_deleted = {} async def record_cache_state_then_delete(*args, **kwargs): - if kwargs.get("table_name") == "team": - cache_state_when_rows_deleted["doomed_still_cached"] = ( - fresh_cache.get_cache(key="team_id:team-doomed") is not None - ) - return {"deleted_teams": ["team-doomed"]} + cache_state_when_rows_deleted["doomed_still_cached"] = ( + fresh_cache.get_cache(key="team_id:team-doomed") is not None + ) + return 1 mock_prisma_client = AsyncMock() mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=doomed_team) - mock_prisma_client.delete_data = AsyncMock(side_effect=record_cache_state_then_delete) + mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_keys": 0}) mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) @@ -7491,6 +7697,7 @@ async def test_delete_team_sweeps_references_outside_members_with_roles( mock_prisma_client.db.execute_raw = mock_execute_raw mock_membership_delete_many = AsyncMock() mock_prisma_client.db.litellm_teammembership.delete_many = mock_membership_delete_many + mock_prisma_client.db.litellm_teamtable.delete_many = AsyncMock(side_effect=record_cache_state_then_delete) mock_tx = AsyncMock() mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) @@ -7499,6 +7706,11 @@ async def test_delete_team_sweeps_references_outside_members_with_roles( mock_tx_cm.__aexit__ = AsyncMock(return_value=False) mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + # The locked delete-and-sweep transaction /team/member_add serializes against, kept + # separate from mock_tx above (the BYOK-model-cleanup transaction, unrelated to this lock). + _wire_team_delete_tx(mock_prisma_client) + mock_lock_tx = mock_prisma_client.tx.return_value.__aenter__.return_value + fresh_cache = UserApiKeyCache() for cached_team_id, cached_alias in ( ("team-doomed", "doomed-team"), @@ -7532,14 +7744,22 @@ async def test_delete_team_sweeps_references_outside_members_with_roles( assert mock_execute_raw.await_args_list == [ call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-doomed"), call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-doomed"), - ], "the sweep must run once before the team row is deleted and again after, so a member_add racing the delete cannot leave the reference behind" + ], ( + "the unlocked sweep must run once to catch pre-existing drift, and the locked sweep " + "(alongside the delete, under the same advisory lock member_add takes) must run again " + "so a member_add that wrote its reference just before losing the lock is still reaped" + ) - # same two passes: the second one reaps a membership row inserted while the delete was running + # same two passes for the membership rows, the second under the lock alongside the delete assert mock_membership_delete_many.await_args_list == [ call(where={"team_id": {"in": ("team-doomed",)}}), call(where={"team_id": {"in": ("team-doomed",)}}), ] + assert mock_lock_tx.query_raw.await_args_list == [call(TEAM_ADVISORY_LOCK_SQL, "team-doomed")], ( + "the advisory lock must be acquired before the team row is deleted" + ) + assert fresh_cache.get_cache(key="team_id:team-doomed") is None assert fresh_cache.get_cache(key="team_alias:doomed-team") is None assert fresh_cache.get_cache(key="team_id:team-kept") is not None @@ -7589,6 +7809,7 @@ async def test_delete_team_evicts_the_auth_cache_of_the_keys_it_deletes( mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) mock_tx_cm.__aexit__ = AsyncMock(return_value=False) mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + _wire_team_delete_tx(mock_prisma_client) fresh_cache = UserApiKeyCache() fresh_cache.set_cache(key="hashed-doomed-key", value=UserAPIKeyAuth(token="hashed-doomed-key", team_id="team-doomed")) @@ -7616,14 +7837,17 @@ async def test_delete_team_evicts_the_auth_cache_of_the_keys_it_deletes( @pytest.mark.asyncio -async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cache( +async def test_delete_team_failing_locked_sweep_rolls_back_the_delete_and_leaves_the_cache_alone( monkeypatch, disable_audit_logging_for_mocked_team, ): """ - The reconcile sweep runs after the team row is committed deleted. If it ran before cache - eviction, a sweep failure would return an error with the team gone from the db but still - served from cache, which is the exact bug this PR exists to fix. + The team delete and its post-delete reconcile sweep run inside one transaction, under the + team's advisory lock, so a sweep failure rolls the delete back with it rather than leaving + the row gone with the sweep half done. Cache eviction only runs after that transaction + commits, so a failure here must leave the team exactly as it was: still in the db, and + still cached. Evicting a cache entry for a delete that never actually committed would be + the same class of bug this PR exists to fix, just on the other side of the transaction. """ from litellm.proxy._types import DeleteTeamRequest from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -7643,7 +7867,7 @@ async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cac mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() - # the first sweep succeeds, the post-delete reconcile sweep blows up + # the unlocked pre-delete sweep succeeds, the locked post-delete sweep blows up mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[None, ConnectionError("db went away")]) mock_tx = AsyncMock() @@ -7652,6 +7876,7 @@ async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cac mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) mock_tx_cm.__aexit__ = AsyncMock(return_value=False) mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + _wire_team_delete_tx(mock_prisma_client) fresh_cache = UserApiKeyCache() cached_obj = LiteLLM_TeamTableCachedObj(team_id="team-doomed", team_alias="doomed-team") @@ -7675,9 +7900,10 @@ async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cac litellm_changed_by="admin-user", ) - # the delete committed, so the cache must not still be serving the team - assert fresh_cache.get_cache(key="team_id:team-doomed") is None - assert fresh_cache.get_cache(key="team_alias:doomed-team") is None + # the transaction that deletes the row and runs the locked sweep never committed, so + # cache eviction (which only runs after that commit) must never have been reached + assert fresh_cache.get_cache(key="team_id:team-doomed") is not None + assert fresh_cache.get_cache(key="team_alias:doomed-team") is not None @pytest.mark.asyncio @@ -7719,6 +7945,7 @@ async def test_delete_team_broadcasts_cache_invalidation_to_other_workers( mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) mock_tx_cm.__aexit__ = AsyncMock(return_value=False) mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + _wire_team_delete_tx(mock_prisma_client) published = [] @@ -7789,6 +8016,7 @@ async def test_delete_team_survives_a_failing_cache_backend( mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) mock_tx_cm.__aexit__ = AsyncMock(return_value=False) mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + _wire_team_delete_tx(mock_prisma_client) exploding_logging_obj = MagicMock() exploding_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock( @@ -7813,7 +8041,7 @@ async def test_delete_team_survives_a_failing_cache_backend( ) assert result == {"deleted_teams": ["team-doomed"]} - mock_delete_data.assert_any_await(team_id_list=["team-doomed"], table_name="team") + mock_prisma_client.db.litellm_teamtable.delete_many.assert_any_await(where={"team_id": {"in": ["team-doomed"]}}) assert exploding_logging_obj.internal_usage_cache.dual_cache.async_delete_cache.await_count > 0 @@ -12603,3 +12831,376 @@ async def test_invalidate_access_group_cache_deletes_the_cached_object(): "user_api_key_cache": cache, "proxy_logging_obj": logging_obj, } + + +def test_validate_team_member_reset_spend_value_rejects_non_numeric(): + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to="not-a-number", + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert exc.value.status_code == 400 + + +def test_validate_team_member_reset_spend_value_rejects_negative(): + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to=-1.0, + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert exc.value.status_code == 400 + + +@pytest.mark.parametrize("reset_to", [float("nan"), float("inf"), float("-inf")]) +def test_validate_team_member_reset_spend_value_rejects_non_finite(reset_to): + """NaN and +/-inf are instances of float and compare False against every bound + below (`nan < 0`, `nan > current_spend` are both False), so an isinstance-and-range + check alone lets them through to persist as the member's spend and silently + disable every later budget comparison against it.""" + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to=reset_to, + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert exc.value.status_code == 400 + + +@pytest.mark.parametrize("reset_to", [True, False]) +def test_reset_spend_request_rejects_bool_reset_to(reset_to): + """bool is a subclass of int, so pydantic silently coerces True/False into 1.0/0.0 for a + ``float`` field: {"reset_to": true} would otherwise reach _validate_team_member_reset_spend_value + as an indistinguishable 1.0 and reset the member's spend instead of failing the request.""" + with pytest.raises(ValidationError): + ResetSpendRequest(reset_to=reset_to) + + +def test_validate_team_member_reset_spend_value_rejects_above_current_spend(): + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to=20.0, + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert exc.value.status_code == 400 + + +def test_validate_team_member_reset_spend_value_rejects_above_max_budget(): + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to=10.0, + membership=LiteLLM_TeamMembership( + user_id="u1", + team_id="t1", + spend=10.0, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=5.0), + ), + ) + assert exc.value.status_code == 400 + + +def test_validate_team_member_reset_spend_value_accepts_valid_reset(): + result = _validate_team_member_reset_spend_value( + reset_to=0.0, + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert result == 0.0 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_success(monkeypatch): + """A proxy admin resetting a stuck team member's spend must write the DB + row to reset_to AND invalidate the cached spend/membership state, or the + 429 the endpoint exists to clear keeps firing off the stale cache. + Asserted against real cache reads, not mock call args, so a change that + keeps the call but drops its effect still fails.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + mock_prisma_client = MagicMock() + mock_proxy_logging_obj = MagicMock() + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership") + await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership") + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0) + + membership_row = LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + spend=10.0, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=50.0), + ) + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")), + ): + response = await reset_team_member_spend_fn( + team_id="team-1", + user_id="member-1", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert response["spend"] == 0.0 + assert response["previous_spend"] == 10.0 + assert response["max_budget"] == 50.0 + mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with( + where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}}, + data={"spend": 0.0}, + ) + assert await real_cache.async_get_cache(key="team-1_member-1") is None + assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 0.0 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_membership_not_found(monkeypatch): + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_spend_fn( + team_id="team-1", + user_id="ghost-user", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_team_not_found(monkeypatch): + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist in db."})), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_spend_fn( + team_id="ghost-team", + user_id="member-1", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_forbidden_for_non_admin(monkeypatch): + """A caller who is neither proxy admin, org admin, nor this team's admin must be refused, + matching every other team-mutating endpoint's authorization.""" + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1", members_with_roles=[])), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_spend_fn( + team_id="team-1", + user_id="member-1", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user", user_id="plain-user" + ), + ) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_team_admin_cannot_reset_own_spend(monkeypatch): + """_verify_team_access authorizes a team admin over their own team with no check that the + target differs from the caller. Unchecked, that admin could target their own membership row + and repeatedly zero it right before it crosses their per-member cap, consuming the shared + team budget without the configured limit ever binding (Veria finding on PR #37971).""" + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + team_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-admin", user_id="team-admin-1") + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock( + return_value=LiteLLM_TeamTable( + team_id="team-1", + members_with_roles=[Member(user_id="team-admin-1", role="admin")], + ) + ), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_spend_fn( + team_id="team-1", + user_id="team-admin-1", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=team_admin, + ) + assert exc.value.status_code == 403 + mock_prisma_client.db.litellm_teammembership.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkeypatch): + """The self-reset guard is scoped to non-proxy-admin roles: a proxy admin resetting their + own membership spend is the platform-wide trust boundary, not a team-scoped one.""" + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + membership_row = LiteLLM_TeamMembership(user_id="admin-user", team_id="team-1", spend=10.0) + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")), + ): + response = await reset_team_member_spend_fn( + team_id="team-1", + user_id="admin-user", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + assert response["spend"] == 0.0 + + +@pytest.mark.asyncio +async def test_team_member_update_invalidates_team_member_spend_state_when_budget_patch_applied(monkeypatch): + """Raising a stuck member's max_budget_in_team via the documented /team/member_update + endpoint must invalidate the cached membership state, or the raised cap never reaches the + admission check and the member stays 429ing. The live spend counter itself must be left + untouched: only the cap changed, and deleting the counter would force a reseed from the + DB's own spend column, which lags the live counter via periodic batch writes, briefly + UNDER-enforcing the raised cap against a spend value lower than what was actually tracked. + Asserted against real cache reads, not mock call args, so a change that keeps the call but + drops its effect still fails.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + mock_prisma_client = MagicMock() + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership") + await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership") + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0) + + team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[]) + team_info_response = { + "team_info": team_row, + "team_memberships": [LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id=None)], + } + + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache) + + mock_tx = AsyncMock() + mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx) + mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None) + + with ( + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.team_info", + AsyncMock(return_value=team_info_response), + ), + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", + AsyncMock(), + ), + ): + await team_member_update( + data=TeamMemberUpdateRequest(team_id="team-1", user_id="member-1", max_budget_in_team=999999.0), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert await real_cache.async_get_cache(key="team-1_member-1") is None + assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 999.0 + + +@pytest.mark.asyncio +async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent(monkeypatch): + """A role-only update carries an empty budget_patch and touches no budget state, + so the member's cached spend/membership state must be left untouched.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + mock_prisma_client = MagicMock() + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_member-1", value="still-fresh-membership") + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=1.5) + + team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[]) + team_info_response = { + "team_info": team_row, + "team_memberships": [LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id=None)], + } + + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache) + + mock_tx = AsyncMock() + mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx) + mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None) + + with ( + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.team_info", + AsyncMock(return_value=team_info_response), + ), + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", + AsyncMock(), + ), + ): + await team_member_update( + data=TeamMemberUpdateRequest(team_id="team-1", user_id="member-1"), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert await real_cache.async_get_cache(key="team-1_member-1") == "still-fresh-membership" + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 1.5 diff --git a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py index bdc2f9065b9..a6b1fc32eda 100644 --- a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py @@ -1120,3 +1120,81 @@ async def test_add_new_member_creates_missing_user_atomically_via_upsert(): assert upsert_data["create"]["teams"] == ["team-1"] assert upsert_data["update"], "empty update branch degrades the upsert to a racy SELECT-then-INSERT" assert "teams" not in upsert_data["update"] + + +def _member_write_tx() -> MagicMock: + tx = MagicMock() + created_user = MagicMock() + created_user.user_id = "pool-user" + created_user.model_dump.return_value = { + "user_id": "pool-user", + "user_email": "pool@example.com", + "teams": ["team-pool"], + "user_role": "internal_user", + } + created_budget = MagicMock() + created_budget.budget_id = "budget-pool" + membership = MagicMock() + membership.model_dump.return_value = { + "team_id": "team-pool", + "user_id": "pool-user", + "budget_id": "budget-pool", + "litellm_budget_table": None, + } + tx.litellm_usertable.upsert = AsyncMock(return_value=created_user) + tx.litellm_usertable.create = AsyncMock(return_value=created_user) + tx.litellm_usertable.update_many = AsyncMock() + tx.litellm_usertable.find_many = AsyncMock(return_value=[]) + tx.litellm_budgettable.find_unique = AsyncMock(return_value=None) + tx.litellm_budgettable.create = AsyncMock(return_value=created_budget) + tx.litellm_teammembership.create = AsyncMock(return_value=membership) + return tx + + +@pytest.mark.parametrize( + "new_member", + [ + Member(user_id="pool-user", role="user"), + Member(user_email="pool@example.com", role="user"), + ], + ids=["by_user_id", "by_user_email"], +) +@pytest.mark.asyncio +async def test_add_new_member_runs_every_write_on_the_caller_transaction(new_member): + """ + Regression pin against exhausting the connection pool with advisory-lock waiters. + + /team/member_add calls this while holding the team's advisory lock inside a transaction, + so it already owns a pooled connection. Any query issued on the regular client here needs + a second one, and enough concurrent adds for one team leave every connection parked on the + lock while the holder waits for a free one, so nothing ever commits or releases the lock. + Given a transaction, every read and write has to go through it. + """ + from litellm.proxy._types import LitellmUserRoles + + tx = _member_write_tx() + prisma_client = AsyncMock() + + result_user, result_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=50.0, + prisma_client=prisma_client, + team_id="team-pool", + user_api_key_dict=UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ), + litellm_proxy_admin_name="admin", + tx=tx, + ) + + assert result_user.user_id == "pool-user" + assert result_membership is not None + assert result_membership.budget_id == "budget-pool" + + assert tx.litellm_budgettable.create.await_count == 1 + assert tx.litellm_teammembership.create.await_count == 1 + assert tx.litellm_usertable.upsert.await_count + tx.litellm_usertable.create.await_count == 1 + + prisma_client.db.assert_not_called() + prisma_client.get_data.assert_not_awaited() + prisma_client.insert_data.assert_not_awaited() diff --git a/tests/test_litellm/proxy/rerank_endpoints/__init__.py b/tests/test_litellm/proxy/rerank_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py b/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py new file mode 100644 index 00000000000..9f11ff6f20d --- /dev/null +++ b/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py @@ -0,0 +1,120 @@ +""" +Tests for rerank_endpoints/endpoints.py response headers. +""" + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import Request, Response + +import litellm.proxy.common_request_processing as common_request_processing_mod +import litellm.proxy.proxy_server as proxy_server_mod +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.rerank_endpoints.endpoints import rerank +from litellm.types.utils import RerankResponse + +HIDDEN_PARAMS = { + "model_id": "deployment-1", + "api_base": "https://bedrock-agent-runtime.us-east-1.amazonaws.com", + "response_cost": 0.002, + "_response_ms": 1500.5, + "litellm_overhead_time_ms": 12.5, + "callback_duration_ms": 1.25, + "timing_llm_api_ms": 1488.0, + "timing_pre_processing_ms": 10.0, + "timing_post_processing_ms": 2.5, + "timing_message_copy_ms": 0.01, +} + + +def _build_request() -> Request: + body = json.dumps({"model": "rerank-model", "query": "q", "documents": ["a", "b"]}).encode() + + async def receive(): + return {"type": "http.request", "body": body, "more_body": False} + + return Request( + scope={ + "type": "http", + "method": "POST", + "path": "/rerank", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + }, + receive=receive, + ) + + +async def _call_rerank(hidden_params: dict = HIDDEN_PARAMS) -> Response: + response = RerankResponse(id="rerank-1", results=[{"index": 0, "relevance_score": 0.9}]) + response._hidden_params = dict(hidden_params) + + fastapi_response = Response() + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + proxy_logging_obj.update_request_status = AsyncMock() + + async def fake_add_litellm_data_to_request(**kwargs): + return {**kwargs["data"], "litellm_call_id": "call-123"} + + async def fake_route_request(**kwargs): + async def _call(): + return response + + return _call() + + with ( + patch.object(proxy_server_mod, "add_litellm_data_to_request", fake_add_litellm_data_to_request), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler + patch.object(proxy_server_mod, "route_request", fake_route_request), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler + patch.object(proxy_server_mod, "proxy_logging_obj", proxy_logging_obj), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler + patch.object(proxy_server_mod, "llm_router", MagicMock()), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler + patch.object(proxy_server_mod, "version", "1.2.3"), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler + ): + await rerank( + request=_build_request(), + fastapi_response=fastapi_response, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + return fastapi_response + + +@pytest.mark.asyncio +async def test_rerank_emits_latency_and_cost_headers(): + """/rerank must surface the same hidden_params-derived headers as /chat/completions.""" + fastapi_response = await _call_rerank() + + assert fastapi_response.headers["x-litellm-call-id"] == "call-123" + assert fastapi_response.headers["x-litellm-response-duration-ms"] == "1500.5" + assert fastapi_response.headers["x-litellm-overhead-duration-ms"] == "12.5" + assert fastapi_response.headers["x-litellm-callback-duration-ms"] == "1.25" + assert fastapi_response.headers["x-litellm-response-cost"] == "0.002" + + +@pytest.mark.asyncio +async def test_rerank_emits_detailed_timing_headers_when_enabled(): + """LITELLM_DETAILED_TIMING must also work on /rerank, not just /chat/completions.""" + with patch.object(common_request_processing_mod, "LITELLM_DETAILED_TIMING", True): # test-quality-ok: LITELLM_DETAILED_TIMING is a module constant; toggling it is the behavior under test + fastapi_response = await _call_rerank() + + assert fastapi_response.headers["x-litellm-timing-llm-api-ms"] == "1488.0" + assert fastapi_response.headers["x-litellm-timing-pre-processing-ms"] == "10.0" + assert fastapi_response.headers["x-litellm-timing-post-processing-ms"] == "2.5" + assert fastapi_response.headers["x-litellm-timing-message-copy-ms"] == "0.01" + + +@pytest.mark.asyncio +async def test_rerank_emits_zero_response_cost_header(): + """A free deployment costs 0.0, which is a real cost and must not be dropped.""" + fastapi_response = await _call_rerank({**HIDDEN_PARAMS, "response_cost": 0.0}) + + assert fastapi_response.headers["x-litellm-response-cost"] == "0.0" + + +@pytest.mark.asyncio +async def test_rerank_omits_detailed_timing_headers_when_disabled(): + with patch.object(common_request_processing_mod, "LITELLM_DETAILED_TIMING", False): # test-quality-ok: LITELLM_DETAILED_TIMING is a module constant; toggling it is the behavior under test + fastapi_response = await _call_rerank() + + assert "x-litellm-timing-llm-api-ms" not in fastapi_response.headers diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 9177944df2d..791d64c6428 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -1353,6 +1353,8 @@ class TestParseCursorModelVariant: ("claude-opus-5-fast", "claude-opus-5", None), ("gpt-5.6-sol", "gpt-5.6-sol", None), ("foo-thinking-ultra-fast", "foo-thinking-ultra", None), + ("gpt-5.6-thinking-max", "gpt-5.6", "max"), + ("foo-thinking-mega-fast", "foo-thinking-mega", None), ("-thinking-high", "-thinking-high", None), ], ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 6c8e641642b..843e1d1296f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -3,11 +3,10 @@ import datetime import json from datetime import timezone from typing import Any, Final, cast +from unittest.mock import AsyncMock, MagicMock, patch import pytest - - -from unittest.mock import AsyncMock, MagicMock, patch +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import ( @@ -3500,6 +3499,50 @@ def test_get_logging_payload_failed_request_without_standard_logging_payload_lea assert payload["custom_llm_provider"] == "" +class _ModelRouterSpendLogKwargs(TypedDict): + model: ReadOnly[str] + litellm_params: ReadOnly[dict[str, dict[str, str]]] + standard_logging_object: ReadOnly[StandardLoggingPayload] + + +def _model_router_spend_log_kwargs(slp_model: str | None) -> _ModelRouterSpendLogKwargs: + standard_logging_payload: Final = cast( + StandardLoggingPayload, + { + "model": slp_model, + "metadata": {}, + "model_map_information": StandardLoggingModelInformation( + model_map_key="azure_ai/model_router", model_map_value=None + ), + }, + ) + return { + "model": "azure_ai/model_router/model-router", + "litellm_params": {"metadata": {"user_api_key": "sk-test-key"}}, + "standard_logging_object": standard_logging_payload, + } + + +def test_get_logging_payload_uses_standard_logging_payload_model(): + payload = get_logging_payload( + kwargs=_model_router_spend_log_kwargs(slp_model="azure_ai/gpt-5-mini"), + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["model"] == "azure_ai/gpt-5-mini" + + +def test_get_logging_payload_falls_back_to_kwargs_model_when_slp_model_missing(): + payload = get_logging_payload( + kwargs=_model_router_spend_log_kwargs(slp_model=None), + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["model"] == "azure_ai/model_router/model-router" + + @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) def test_get_logging_payload_empty_key_slp_none_is_empty_string_not_none_literal(): diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 58714a5e319..b595c44d2ce 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2294,6 +2294,54 @@ class TestOverrideOpenAIResponseModel: assert response_obj.model == actual_model_used assert response_obj.model != requested_model + def test_override_model_preserves_model_router_model_for_alias_without_router_in_name( + self, + ): + """ + The client sends a model group alias, which carries no model_router/ prefix, so the + name check alone only fires when the operator happened to put "model-router" in the + alias. With the stamp on the response the actual model survives whatever it is named. + """ + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + ) + + requested_model = "smart-pick" + actual_model_used = "azure_ai/grok-4-1-fast-reasoning" + + response_obj = MagicMock() + response_obj.model = actual_model_used + response_obj._hidden_params = { + "additional_headers": {}, + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: actual_model_used, + } + + _override_openai_response_model( + response_obj=response_obj, + requested_model=requested_model, + log_context="test_context", + ) + assert response_obj.model == actual_model_used + + def test_override_model_still_restamps_non_router_alias_without_stamp(self): + """ + Control for the test above: absent the stamp, an ordinary deployment keeps being + restamped to the requested model, so the stamp is doing the work rather than the + preserve branch having gone unconditional. + """ + requested_model = "smart-pick" + + response_obj = MagicMock() + response_obj.model = "azure_ai/grok-4-1-fast-reasoning" + response_obj._hidden_params = {"additional_headers": {}} + + _override_openai_response_model( + response_obj=response_obj, + requested_model=requested_model, + log_context="test_context", + ) + assert response_obj.model == requested_model + def test_override_model_uses_winning_model_for_fastest_response(self): """ Test that when fastest_response batch completion is used with a @@ -7295,3 +7343,31 @@ class TestRouterModelNameOnNonStreamingResponse: ) assert response[ROUTER_MODEL_NAME_RESPONSE_FIELD] == "deep-model" + + +@pytest.mark.parametrize( + "exc,expect_traceback", + [ + pytest.param(HTTPException(status_code=400, detail="Invalid model name passed in"), False, id="expected_400"), + pytest.param(ValueError("unexpected internal error"), True, id="unexpected_error"), + ], +) +def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_traceback, caplog): + """Regression for LIT-6043: expected 4xx errors log without formatting a + traceback; unexpected errors keep logger.exception behavior.""" + from litellm._logging import verbose_proxy_logger + from litellm.proxy.common_request_processing import _log_llm_api_exception + + verbose_proxy_logger.propagate = True + try: + with caplog.at_level("ERROR", logger="LiteLLM Proxy"): + try: + raise exc + except Exception as raised: + _log_llm_api_exception(raised) + finally: + verbose_proxy_logger.propagate = False + + records = [r for r in caplog.records if "_handle_llm_api_exception(): Exception occured" in r.getMessage()] + assert len(records) == 1 + assert (records[0].exc_info is not None) is expect_traceback diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 31d2a6cef98..b9a31acca96 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -78,6 +78,16 @@ def client_no_auth(): return TestClient(app) +def test_cors_exposes_cache_key_header_to_browser_js(): + from fastapi.middleware.cors import CORSMiddleware + + from litellm.constants import LITELLM_UI_ALLOW_HEADERS + + cors_middleware = next(m for m in app.user_middleware if m.cls is CORSMiddleware) + assert cors_middleware.kwargs["expose_headers"] is LITELLM_UI_ALLOW_HEADERS + assert "x-litellm-cache-key" in cors_middleware.kwargs["expose_headers"] + + def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): mock_login_result = {"user_id": "test-user"} mock_prisma_client = MagicMock() @@ -11347,3 +11357,38 @@ class TestRouterModelNameOnStreamingChunks: assert len(frames) >= 3 assert '"router_model_name":"deep-model"' in frames[0] assert all('"router_model_name":"backup-tier"' in frame for frame in frames[1:]) + + +@pytest.mark.asyncio +async def test_authoritative_floor_spend_keeps_a_reset_marker_written_during_the_db_read(): + """A team-member spend reset writes the post-reset floor to the spend_db_floor marker + (auth_checks.invalidate_team_member_spend_state). A floor read already in flight when the + reset commits would otherwise cache its stale pre-reset DB value over the fresh marker, + letting a budget check raise the counter right back above the just-reset spend + (regression: PR #37971 Greptile finding).""" + from litellm.proxy.proxy_server import _authoritative_floor_spend + + real_spend_counter_cache = DualCache() + counter_key = "spend:team_member:user-1:team-1" + marker_key = f"spend_db_floor:{counter_key}" + + async def db_read_racing_with_a_reset(prisma_client, counter_key): + real_spend_counter_cache.in_memory_cache.set_cache(key=marker_key, value=0.0) + return 999.0 + + with ( + patch.object( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + proxy_server_module, "spend_counter_cache", real_spend_counter_cache + ), + patch.object( # test-quality-ok: the DB read must race the reset; no injectable seam for module-global prisma reads + proxy_server_module.SpendCounterReseed, + "from_db", + AsyncMock(side_effect=db_read_racing_with_a_reset), + ), + ): + result = await _authoritative_floor_spend(counter_key=counter_key) + + assert result == 0.0 + assert real_spend_counter_cache.in_memory_cache.get_cache(key=marker_key) == 0.0, ( + "the in-flight DB read clobbered the post-reset floor marker with the stale pre-reset value" + ) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index fb01216982f..6920cc0dae3 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -11,7 +11,6 @@ from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks - from unittest.mock import MagicMock, patch from litellm.proxy.utils import get_custom_url, join_paths @@ -82,9 +81,7 @@ async def test_proxy_only_error_log_marks_no_upstream_llm_call(): captured = {} def fake_pre_call(self, *args, **kwargs): - captured["flag"] = self.model_call_details.get( - LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL - ) + captured["flag"] = self.model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL) from litellm.litellm_core_utils.litellm_logging import Logging @@ -102,9 +99,7 @@ async def test_proxy_only_error_log_marks_no_upstream_llm_call(): "model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], }, - user_api_key_dict=UserAPIKeyAuth( - api_key="sk-bad", request_route="/v1/chat/completions" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-bad", request_route="/v1/chat/completions"), route="/v1/chat/completions", original_exception=Exception("bad key"), ) @@ -148,13 +143,9 @@ async def test_proxy_only_error_log_keeps_litellm_metadata_in_litellm_params(): request_data={ "model": "gpt-4o", "input": "blocked prompt", - "litellm_metadata": { - "standard_logging_guardrail_information": guardrail_info - }, + "litellm_metadata": {"standard_logging_guardrail_information": guardrail_info}, }, - user_api_key_dict=UserAPIKeyAuth( - api_key="sk-1234", request_route="/v1/responses" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/responses"), route="/v1/responses", original_exception=HTTPException(status_code=400, detail="blocked"), ) @@ -163,12 +154,7 @@ async def test_proxy_only_error_log_keeps_litellm_metadata_in_litellm_params(): Logging.pre_call = orig_pre_call Logging.async_failure_handler = orig_async_failure - assert ( - captured["litellm_params"]["litellm_metadata"][ - "standard_logging_guardrail_information" - ] - == guardrail_info - ) + assert captured["litellm_params"]["litellm_metadata"]["standard_logging_guardrail_information"] == guardrail_info assert "litellm_metadata" not in captured["optional_params"] @@ -206,9 +192,7 @@ def test_get_model_group_info_order(): def test_join_paths_no_duplication(): """Test that join_paths doesn't duplicate route when base_path already ends with it""" - result = join_paths( - base_path="http://0.0.0.0:4000/my-custom-path/", route="/my-custom-path" - ) + result = join_paths(base_path="http://0.0.0.0:4000/my-custom-path/", route="/my-custom-path") assert result == "http://0.0.0.0:4000/my-custom-path" @@ -814,9 +798,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: estimated = request_data["combined_usage_object"] assert isinstance(estimated, Usage) - expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( - model="gpt-3.5-turbo", text=system_prompt - ) + expected = litellm_module.token_counter( + model="gpt-3.5-turbo", messages=messages + ) + litellm_module.token_counter(model="gpt-3.5-turbo", text=system_prompt) assert estimated.prompt_tokens == expected @pytest.mark.asyncio @@ -834,7 +818,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: estimated = request_data["combined_usage_object"] assert isinstance(estimated, Usage) - expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( + expected = litellm_module.token_counter( + model="gpt-3.5-turbo", messages=messages + ) + litellm_module.token_counter( model="gpt-3.5-turbo", text="part one of the system prompt. part two of the system prompt." ) assert estimated.prompt_tokens == expected @@ -870,9 +856,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: estimated = request_data["combined_usage_object"] assert isinstance(estimated, Usage) - expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( - model="gpt-3.5-turbo", text=system_prompt - ) + expected = litellm_module.token_counter( + model="gpt-3.5-turbo", messages=messages + ) + litellm_module.token_counter(model="gpt-3.5-turbo", text=system_prompt) assert estimated.prompt_tokens == expected @pytest.mark.asyncio @@ -890,9 +876,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: estimated = request_data["combined_usage_object"] assert isinstance(estimated, Usage) - expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( - model="gpt-3.5-turbo", text=dispatched_system - ) + expected = litellm_module.token_counter( + model="gpt-3.5-turbo", messages=messages + ) + litellm_module.token_counter(model="gpt-3.5-turbo", text=dispatched_system) assert estimated.prompt_tokens == expected @@ -916,9 +902,7 @@ def test_create_model_info_response_includes_max_tokens_from_lookup(): model_id="some-model", provider="openai", llm_router=None, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=128000, max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=128000, max_output_tokens=16384), ) assert response["id"] == "some-model" @@ -935,9 +919,7 @@ def test_create_model_info_response_does_not_call_router_group_info(): model_id="some-model", provider="openai", llm_router=router, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=128000, max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=128000, max_output_tokens=16384), ) router.get_model_group_info.assert_not_called() @@ -968,9 +950,7 @@ def test_create_model_info_response_deployment_limits_override_cost_map(): model_id="gpt-4o", provider="openai", llm_router=router, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=128000, max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=128000, max_output_tokens=16384), ) assert response["max_input_tokens"] == 200000 @@ -1008,9 +988,7 @@ def test_create_model_info_response_survives_malformed_cost_map_limits(bad_value model_id="some-model", provider="openai", llm_router=None, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=bad_value, max_output_tokens=bad_value - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=bad_value, max_output_tokens=bad_value), ) assert response["id"] == "some-model" @@ -1023,9 +1001,7 @@ def test_create_model_info_response_keeps_valid_cost_map_limit_beside_malformed_ model_id="some-model", provider="openai", llm_router=None, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens="128,000", max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens="128,000", max_output_tokens=16384), ) assert "max_input_tokens" not in response @@ -1068,9 +1044,7 @@ def test_create_model_info_response_emits_integer_token_counts(): model_id="some-model", provider="openai", llm_router=None, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=128000, max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=128000, max_output_tokens=16384), ) assert isinstance(response["max_input_tokens"], int) @@ -1119,9 +1093,7 @@ def test_create_model_info_response_no_router_keeps_base_fields(): def test_create_model_info_response_reads_real_cost_map(): - response = create_model_info_response( - model_id="gpt-4o", provider="openai", llm_router=None - ) + response = create_model_info_response(model_id="gpt-4o", provider="openai", llm_router=None) assert isinstance(response["max_input_tokens"], int) assert response["max_input_tokens"] > 0 @@ -1205,10 +1177,7 @@ class TestPostCallFailureHookLLMExceptionAlerting: @pytest.mark.asyncio async def test_http_exception_does_not_alert(self): - assert ( - await self._alerted(HTTPException(status_code=400, detail="blocked")) - is False - ) + assert await self._alerted(HTTPException(status_code=400, detail="blocked")) is False @pytest.mark.asyncio async def test_genuine_llm_api_error_still_alerts(self): @@ -1241,9 +1210,7 @@ class TestPostCallFailureHookProxyExceptionLogging: await proxy_logging_obj.post_call_failure_hook( request_data={}, original_exception=exc, - user_api_key_dict=UserAPIKeyAuth( - api_key="sk-test", request_route=request_route - ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", request_route=request_route), ) return handle_mock.await_count > 0 @@ -1260,20 +1227,12 @@ class TestPostCallFailureHookProxyExceptionLogging: @pytest.mark.asyncio async def test_proxy_exception_on_llm_route_is_logged(self): - assert ( - await self._logged(self._block(), request_route="/v1/chat/completions") - is True - ) + assert await self._logged(self._block(), request_route="/v1/chat/completions") is True @pytest.mark.asyncio async def test_generic_exception_on_llm_route_is_not_logged(self): # A raw provider/unknown exception is logged by the LLM call path, not here. - assert ( - await self._logged( - Exception("upstream 503"), request_route="/v1/chat/completions" - ) - is False - ) + assert await self._logged(Exception("upstream 503"), request_route="/v1/chat/completions") is False class TestShouldUseSmtpSsl: @@ -1307,9 +1266,7 @@ class TestCreateSmtpConnection: patch("smtplib.SMTP_SSL") as mock_smtp_ssl, patch("smtplib.SMTP") as mock_smtp, ): - result = _create_smtp_connection( - smtp_host="mail.example.com", smtp_port=465 - ) + result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=465) mock_smtp.assert_not_called() assert result is mock_smtp_ssl.return_value @@ -1329,9 +1286,7 @@ class TestCreateSmtpConnection: patch("smtplib.SMTP_SSL") as mock_smtp_ssl, patch("smtplib.SMTP") as mock_smtp, ): - result = _create_smtp_connection( - smtp_host="mail.example.com", smtp_port=587 - ) + result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=587) mock_smtp_ssl.assert_not_called() assert result is mock_smtp.return_value @@ -1352,9 +1307,7 @@ class TestSendEmailStartTls: monkeypatch.delenv("SMTP_USE_SSL", raising=False) mock_server = MagicMock(spec=smtplib.SMTP) - with patch( - "litellm.proxy.utils._create_smtp_connection" - ) as mock_create_connection: + with patch("litellm.proxy.utils._create_smtp_connection") as mock_create_connection: mock_create_connection.return_value.__enter__.return_value = mock_server await send_email( receiver_email="receiver@example.com", @@ -1690,9 +1643,7 @@ def test_a_failed_dispatch_is_estimated_as_input_only(): usage = _estimate_dispatched_failure_usage(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) assert usage is not None - assert usage.prompt_tokens == _count_request_input_tokens( - FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None - ) + assert usage.prompt_tokens == _count_request_input_tokens(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) assert usage.completion_tokens == 0 assert usage.total_tokens == usage.prompt_tokens @@ -1759,9 +1710,7 @@ def test_a_request_that_reached_a_provider_bills_its_input_at_no_cost(): def test_a_failure_that_cost_the_provider_nothing_lifts_nothing(model_call_details, dispatched): from litellm.proxy.utils import _failure_usage_to_lift - assert _failure_usage_to_lift( - model_call_details=model_call_details, request_body={}, dispatched=dispatched - ) is None + assert _failure_usage_to_lift(model_call_details=model_call_details, request_body={}, dispatched=dispatched) is None def test_the_no_upstream_call_key_the_module_uses_is_the_one_asserted_above(): @@ -1829,3 +1778,102 @@ def test_a_dispatched_failure_lifts_the_four_fields_the_spend_log_needs(): assert lifted["response_cost"] == 0.0 assert lifted["combined_usage_object"].prompt_tokens > 0 assert lifted["standard_logging_object"] == {"id": "log-1"} + + +@pytest.mark.asyncio +async def test_proxy_only_error_expected_4xx_skips_traceback_for_both_handlers(monkeypatch): + """Regression for LIT-6043: an expected 4xx must not format a traceback for + either the async or the threaded sync failure handler.""" + import asyncio + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy._types import UserAPIKeyAuth + + monkeypatch.setattr(litellm, "failure_callback", []) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + captured = {} + sync_ran = asyncio.Event() + loop = asyncio.get_running_loop() + + async def fake_async_failure(self, exception, traceback_exception, *args, **kwargs): + captured["async_traceback"] = traceback_exception + + def fake_sync_failure(self, exception, traceback_exception, *args, **kwargs): + captured["sync_traceback"] = traceback_exception + loop.call_soon_threadsafe(sync_ran.set) + + orig_async_failure = Logging.async_failure_handler + orig_sync_failure = Logging.failure_handler + Logging.async_failure_handler = fake_async_failure + Logging.failure_handler = fake_sync_failure + try: + try: + raise HTTPException(status_code=400, detail="Invalid model name passed in") + except HTTPException as exc: + await proxy_logging_obj._handle_logging_proxy_only_error( + request_data={ + "model": "does-not-exist", + "messages": [{"role": "user", "content": "hi"}], + }, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), + route="/v1/chat/completions", + original_exception=exc, + ) + await asyncio.wait_for(sync_ran.wait(), timeout=5) + finally: + Logging.async_failure_handler = orig_async_failure + Logging.failure_handler = orig_sync_failure + + assert captured["async_traceback"] == "" + assert captured["sync_traceback"] == "" + + +@pytest.mark.asyncio +async def test_proxy_only_error_5xx_keeps_traceback_and_runs_sync_callbacks(monkeypatch): + """Unexpected (5xx) errors keep the full traceback, and a configured + sync-only failure callback still gets its threaded handler.""" + import asyncio + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy._types import UserAPIKeyAuth + + def _custom_sync_callback(kwargs, completion_response, start_time, end_time): + pass + + monkeypatch.setattr(litellm, "failure_callback", [_custom_sync_callback]) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + captured = {} + sync_ran = asyncio.Event() + loop = asyncio.get_running_loop() + + async def fake_async_failure(self, exception, traceback_exception, *args, **kwargs): + captured["async_traceback"] = traceback_exception + + def fake_sync_failure(self, *args, **kwargs): + loop.call_soon_threadsafe(sync_ran.set) + + orig_async_failure = Logging.async_failure_handler + orig_sync_failure = Logging.failure_handler + Logging.async_failure_handler = fake_async_failure + Logging.failure_handler = fake_sync_failure + try: + try: + raise HTTPException(status_code=500, detail="internal error") + except HTTPException as exc: + await proxy_logging_obj._handle_logging_proxy_only_error( + request_data={ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + }, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), + route="/v1/chat/completions", + original_exception=exc, + ) + await asyncio.wait_for(sync_ran.wait(), timeout=5) + finally: + Logging.async_failure_handler = orig_async_failure + Logging.failure_handler = orig_sync_failure + + assert "test_proxy_utils" in captured["async_traceback"] diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index 38af52f165c..758d379f22c 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -541,17 +541,19 @@ class TestTeamRepository: assert [m.user_id for m in members] == expected_ids sql = tx.query_raw.call_args.args[0] - assert "FOR UPDATE" in sql + assert "FOR UPDATE" not in sql, ( + "a row lock here can deadlock with the access-group endpoints; the caller must " + "already hold the team's advisory lock, so a plain read is all this needs" + ) assert tx.query_raw.call_args.args[1] == "team-1" @pytest.mark.asyncio async def test_get_members_with_roles_locked_missing_row(self, repo): """None, not [], so a caller can tell a deleted team from an empty one. - /team/member_add reconciles membership under this lock and has to fail, - and clean up the references it already wrote, when a /team/delete - committed underneath it. An empty list would look like a live team with - no members and it would carry on writing. + /team/member_add reconciles membership under the team's advisory lock and has to + fail, without writing anything, when a /team/delete committed underneath it. An + empty list would look like a live team with no members and it would carry on writing. """ tx = MagicMock() tx.query_raw = AsyncMock(return_value=[]) diff --git a/tests/test_litellm/rerank_api/test_main.py b/tests/test_litellm/rerank_api/test_main.py index 85777afe81c..587be59c550 100644 --- a/tests/test_litellm/rerank_api/test_main.py +++ b/tests/test_litellm/rerank_api/test_main.py @@ -1,6 +1,10 @@ import logging from unittest.mock import MagicMock, patch +import httpx +import pytest +import respx + import litellm @@ -62,3 +66,66 @@ def test_rerank_does_not_log_request_content_at_info(caplog): assert all( r.levelno == logging.DEBUG for r in optional_params_logs ), "optional_rerank_params must be logged at DEBUG, not INFO" + + +TOGETHER_RERANK_BODY = { + "id": "rerank-mock-id", + "results": [{"index": 0, "relevance_score": 0.95}], + "usage": {"prompt_tokens": 10, "total_tokens": 10}, +} + + +def test_together_rerank_defaults_to_together_ai_host(respx_mock: respx.MockRouter, monkeypatch): + """Regression for the Together host migration: rerank used to hardcode + https://api.together.xyz/v1/rerank. The default must now be api.together.ai.""" + monkeypatch.delenv("TOGETHER_AI_API_BASE", raising=False) + + mock_route = respx_mock.post("https://api.together.ai/v1/rerank") + mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY) + + response = litellm.rerank( + model="together_ai/mixedbread-ai/mxbai-rerank-large-v2", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-together-key", + ) + + assert mock_route.called + assert response.results[0]["relevance_score"] == 0.95 + + +def test_together_rerank_honors_api_base(respx_mock: respx.MockRouter): + """Regression: a custom api_base was silently ignored by the Together rerank handler.""" + mock_route = respx_mock.post("https://custom-together.example/v1/rerank") + mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY) + + litellm.rerank( + model="together_ai/mixedbread-ai/mxbai-rerank-large-v2", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-together-key", + api_base="https://custom-together.example/v1", + ) + + assert mock_route.called + assert mock_route.calls[0].request.headers["authorization"] == "Bearer fake-together-key" + + +@pytest.mark.asyncio +async def test_together_rerank_async_honors_env_api_base(respx_mock: respx.MockRouter, monkeypatch): + """Regression: TOGETHER_AI_API_BASE was honored by chat but ignored by rerank.""" + monkeypatch.setenv("TOGETHER_AI_API_BASE", "https://env-together.example/v1") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + + mock_route = respx_mock.post("https://env-together.example/v1/rerank") + mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY) + + response = await litellm.arerank( + model="together_ai/mixedbread-ai/mxbai-rerank-large-v2", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-together-key", + ) + + assert mock_route.called + assert response.results[0]["relevance_score"] == 0.95 diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py new file mode 100644 index 00000000000..6cd5b956b05 --- /dev/null +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -0,0 +1,198 @@ +import pytest + +from litellm.router_utils.reasoning_effort_capability import ( + deployment_is_catalog_mapped, + intersect_supported_reasoning_efforts, + resolve_supported_reasoning_efforts, +) + + +class TestDeploymentIsCatalogMapped: + def test_a_mode_the_catalog_supplied_marks_the_deployment_mapped(self): + assert deployment_is_catalog_mapped({"mode": "chat"}, {}) is True + + def test_a_deployment_the_catalog_never_described_is_not_mapped(self): + assert deployment_is_catalog_mapped(None, {}) is False + assert deployment_is_catalog_mapped({"max_input_tokens": 200000}, {}) is False + + def test_a_mode_the_operator_wrote_does_not_make_the_deployment_mapped(self): + # Every deployment is registered in the cost map under its own id, so an operator-written + # mode reads back identically to one the catalog supplied and would otherwise let an + # off-map deployment empty the levels its mapped siblings agree on. + assert deployment_is_catalog_mapped({"mode": "chat"}, {"mode": "chat", "id": "abc"}) is False + + +class TestProvenanceSeparatesUnknownFromNonReasoning: + def test_an_off_map_deployment_resolves_to_unknown(self): + # get_model_info answers supports_reasoning None both for a deployment the map never + # described and for a mapped non-reasoning model, so reading an unset flag as () would let + # one custom deployment empty every level its mapped siblings agree on. + assert resolve_supported_reasoning_efforts({}, deployment_is_mapped=False) is None + assert resolve_supported_reasoning_efforts({"supports_reasoning": None}, deployment_is_mapped=False) is None + + def test_a_mapped_deployment_the_map_calls_non_reasoning_supports_no_efforts(self): + assert resolve_supported_reasoning_efforts({}, deployment_is_mapped=True) == () + assert resolve_supported_reasoning_efforts({"supports_reasoning": None}, deployment_is_mapped=True) == () + + def test_an_explicit_false_supports_no_efforts_off_the_map_too(self): + # The operator's own escape hatch: saying so on an off-map deployment must still empty the + # group, since nothing else can tell the resolver that model takes no effort level. + assert resolve_supported_reasoning_efforts({"supports_reasoning": False}, deployment_is_mapped=False) == () + + +class TestResolveSupportedReasoningEfforts: + def test_a_reasoning_model_with_no_flags_at_all_resolves_to_unknown(self): + # 689 of the map's 854 reasoning entries carry no effort flag, and the o-series, xai and + # bedrock nova entries among them accept neither none nor minimal, so composing a set out of + # the opt-out defaults alone would advertise levels those providers reject. + assert resolve_supported_reasoning_efforts({"supports_reasoning": True}, deployment_is_mapped=True) is None + + def test_explicit_false_removes_an_opt_out_level(self): + # The gpt-5.5-pro shape from the model map: only medium/high/xhigh are accepted upstream. + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "supports_none_reasoning_effort": False, + "supports_minimal_reasoning_effort": False, + "supports_low_reasoning_effort": False, + "supports_xhigh_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved == ("medium", "high", "xhigh") + + def test_explicit_true_adds_the_opt_in_levels(self): + # The claude-opus shape: xhigh and max explicitly true, everything else absent. + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "supports_xhigh_reasoning_effort": True, + "supports_max_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved == ("none", "minimal", "low", "medium", "high", "xhigh", "max") + + def test_opt_in_flag_set_false_stays_excluded(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "supports_minimal_reasoning_effort": True, + "supports_xhigh_reasoning_effort": False, + }, + deployment_is_mapped=True, + ) + assert resolved == ("none", "minimal", "low", "medium", "high") + + +class TestBareModelNameFallback: + def test_a_prefixed_entry_inherits_the_flags_of_its_unprefixed_twin(self): + """azure/gpt-5-mini carries no effort flag while gpt-5-mini carries three, and the request + path resolves capability flags through that same twin (#20885). Reading only the prefixed + entry would answer unknown for a model the map fully describes.""" + from litellm.utils import _get_model_info_helper + + model_info = dict(_get_model_info_helper(model="gpt-5-mini", custom_llm_provider="azure")) + + assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( + "minimal", + "low", + "medium", + "high", + ) + + def test_the_prefixed_entry_wins_over_its_twin_per_flag(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "azure", + "key": "azure/gpt-5-mini", + "supports_xhigh_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved == ("minimal", "low", "medium", "high", "xhigh") + + +class TestNoneLevelPolarity: + def test_none_stays_opt_out_off_azure(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "openai", + "key": "openai/some-reasoner", + "supports_max_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved is not None and "none" in resolved + + def test_none_stays_opt_out_on_an_azure_model_outside_the_gpt_5_family(self): + """AzureOpenAIGPT5Config is selected by is_model_gpt_5_model, so an azure o-series or + anthropic deployment never reaches the gate that refuses none and must keep the level.""" + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "azure", + "key": "azure/o3", + "supports_max_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved is not None and "none" in resolved + + def test_azure_gpt_5_without_the_flag_does_not_advertise_none(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "azure", + "key": "azure/gpt-5-turbo", + "supports_minimal_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved == ("minimal", "low", "medium", "high") + + def test_azure_gpt_5_with_the_flag_advertises_none(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "azure", + "key": "azure/gpt-5-turbo", + "supports_none_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved is not None and "none" in resolved + + @pytest.mark.parametrize( + "model_key", + ["azure/gpt-5", "azure/gpt-5-mini", "azure/gpt-5-nano", "azure/gpt-5.2", "azure/gpt-5.6"], + ) + def test_azure_advertisement_matches_the_azure_request_gate(self, model_key): + """AzureOpenAIGPT5Config raises UnsupportedParamsError on reasoning_effort='none' for models + it does not flag, so advertising the level there would offer routing a 400.""" + from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config + from litellm.utils import _get_model_info_helper + + model_info = dict(_get_model_info_helper(model=model_key.split("/", 1)[1], custom_llm_provider="azure")) + resolved = resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) + + assert resolved is not None + gate_accepts_none = AzureOpenAIGPT5Config._supports_reasoning_effort_level(model_key, "none") + assert ("none" in resolved) is gate_accepts_none + + +class TestIntersectSupportedReasoningEfforts: + def test_unknown_never_narrows(self): + assert intersect_supported_reasoning_efforts(["medium", "high"], None) == ("medium", "high") + assert intersect_supported_reasoning_efforts(None, ["medium", "high"]) == ("medium", "high") + assert intersect_supported_reasoning_efforts(None, None) is None + + def test_intersection_keeps_canonical_order(self): + assert intersect_supported_reasoning_efforts( + ["max", "high", "medium", "xhigh"], ["xhigh", "medium", "minimal"] + ) == ("medium", "xhigh") + + def test_disjoint_sets_intersect_to_empty(self): + assert intersect_supported_reasoning_efforts(["max"], ["minimal"]) == () diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 99b1cc826aa..8f2b06be4b3 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2944,3 +2944,16 @@ def test_a_stream_that_reported_no_usage_is_still_billed(local_cost_map): assert cost == pytest.approx( _priced_at(rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens) ) + + +@pytest.mark.asyncio +async def test_acompletion_resolves_provider_from_api_base(): + response = await litellm.acompletion( + model="deepseek-chat", + api_base="https://api.deepseek.com/v1", + api_key="fake-key", + messages=[{"role": "user", "content": "hi"}], + mock_response="resolved", + ) + + assert response.choices[0].message.content == "resolved" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 910b874c2ac..56fd6df446f 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -12,8 +12,17 @@ import pytest import litellm +from litellm import Router from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger +from litellm.router import ( + MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS, + FallbackAwareAnthropicMessagesStream, + _anthropic_stream_commits_now, + _anthropic_stream_should_decline_fallback, + _anthropic_stream_should_drop_pre_content_ping, + _is_retriable_anthropic_status, +) def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -8921,3 +8930,1398 @@ class TestAzureBaseModelFallbackLogging: deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] + +def test_model_group_info_intersects_supported_reasoning_efforts(): + router = litellm.Router( + model_list=[ + { + "model_name": "smart-group", + "litellm_params": {"model": "anthropic/opus-like"}, + "model_info": {"id": "opus-like-deployment"}, + }, + { + "model_name": "smart-group", + "litellm_params": {"model": "openai/mini-like"}, + "model_info": {"id": "mini-like-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + if model_id == "opus-like-deployment": + return { + "key": model_name, + "litellm_provider": "anthropic", + "mode": "chat", + "supports_reasoning": True, + "supports_xhigh_reasoning_effort": True, + "supports_max_reasoning_effort": True, + } + return { + "key": model_name, + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "supports_none_reasoning_effort": False, + "supports_minimal_reasoning_effort": True, + "supports_xhigh_reasoning_effort": False, + } + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="smart-group", + user_facing_model_group_name="smart-group", + ) + + assert result is not None + # opus-like offers all seven levels, mini-like lacks none/xhigh/max; only the common set survives, + # so the group never advertises an effort routing could hand to a deployment that rejects it. + assert result.supported_reasoning_efforts == ("minimal", "low", "medium", "high") + + +def test_model_group_info_reasoning_efforts_ignore_a_deployment_off_the_map(): + """The router fills every ModelInfo key, so a deployment absent from the model map arrives with + supports_reasoning None rather than with the key missing. Its synthesized entry carries no mode, + which is what separates it from a mapped non-reasoning model, and nothing being known about it is + no reason to drop the levels the rest of the group agrees on.""" + router = litellm.Router( + model_list=[ + { + "model_name": "smart-group", + "litellm_params": {"model": "anthropic/opus-like"}, + "model_info": {"id": "opus-like-deployment"}, + }, + { + "model_name": "smart-group", + "litellm_params": {"model": "openai/unmapped-model"}, + "model_info": {"id": "unmapped-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + if model_id == "opus-like-deployment": + return { + "key": model_name, + "litellm_provider": "anthropic", + "mode": "chat", + "supports_reasoning": True, + "supports_max_reasoning_effort": True, + } + return {"key": model_name, "litellm_provider": "openai", "mode": None, "supports_reasoning": None} + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="smart-group", + user_facing_model_group_name="smart-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high", "max") + + +def test_model_group_info_reasoning_efforts_empty_on_a_mapped_non_reasoning_deployment(): + """A group mixing a reasoning model with one the map knows is not a reasoning model shares no + level, so it advertises none and the picker offers nothing rather than a level routing would + hand to a deployment that rejects it.""" + router = litellm.Router( + model_list=[ + { + "model_name": "mixed-group", + "litellm_params": {"model": "anthropic/opus-like"}, + "model_info": {"id": "opus-like-deployment"}, + }, + { + "model_name": "mixed-group", + "litellm_params": {"model": "openai/plain-chat"}, + "model_info": {"id": "plain-chat-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + if model_id == "opus-like-deployment": + return { + "key": model_name, + "litellm_provider": "anthropic", + "mode": "chat", + "supports_reasoning": True, + "supports_max_reasoning_effort": True, + } + return {"key": model_name, "litellm_provider": "openai", "mode": "chat", "supports_reasoning": None} + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="mixed-group", + user_facing_model_group_name="mixed-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == () + + +def test_model_group_info_reasoning_efforts_ignore_a_value_declared_in_model_info(): + """The group's levels are computed from its deployments, so a value an operator left in one + deployment's model_info must not seed them. Seeding let the first deployment read narrow the + whole group while the same value on any other deployment was silently ignored.""" + router = litellm.Router( + model_list=[ + { + "model_name": "declared-group", + "litellm_params": {"model": "openai/first-reasoner"}, + "model_info": {"id": "first-deployment"}, + }, + { + "model_name": "declared-group", + "litellm_params": {"model": "openai/second-reasoner"}, + "model_info": {"id": "second-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + info = { + "key": model_name, + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "supports_none_reasoning_effort": True, + } + if model_id == "first-deployment": + info["supported_reasoning_efforts"] = ("high",) + return info + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="declared-group", + user_facing_model_group_name="declared-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high") + + +def test_model_group_info_survives_a_junk_typed_operator_effort_value(): + """A deployment's registered model_info reads back with whatever the operator wrote under any + key, so a wrong-typed supported_reasoning_efforts must not fail the group's info. Only the + constructor's trailing override keeps the junk away from ModelGroupInfo validation.""" + router = litellm.Router( + model_list=[ + { + "model_name": "junk-declared-group", + "litellm_params": {"model": "openai/lone-reasoner"}, + "model_info": {"id": "junk-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + return { + "key": model_name, + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "supports_none_reasoning_effort": True, + "supported_reasoning_efforts": "high", + } + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="junk-declared-group", + user_facing_model_group_name="junk-declared-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high") + + +def test_model_group_info_reasoning_efforts_ignore_a_mode_the_operator_declared(): + """A deployment is registered in the cost map under its own id with whatever model_info the + operator wrote, so a mode they set themselves reads back exactly like one the map supplied. Only + a mode the map supplied marks the deployment as known, or an off-map deployment carrying any + mode empties the group it sits in.""" + from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts + + mapped_model = "openai/gpt-5.6-sol" + expected = resolve_supported_reasoning_efforts( + litellm.get_model_info(model=mapped_model), + deployment_is_mapped=True, + ) + assert expected + + router = litellm.Router( + model_list=[ + { + "model_name": "smart-group", + "litellm_params": {"model": mapped_model, "api_key": "sk-fake"}, + "model_info": {"id": "mapped-deployment"}, + }, + { + "model_name": "smart-group", + "litellm_params": {"model": "openai/a-model-the-map-never-heard-of", "api_key": "sk-fake"}, + "model_info": {"id": "off-map-deployment", "mode": "chat"}, + }, + ] + ) + + result = router._set_model_group_info( + model_group="smart-group", + user_facing_model_group_name="smart-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == expected + + +class TestAddDeploymentApiBaseProviderResolution: + def test_bare_model_with_known_api_base_initializes(self): + router = litellm.Router( + model_list=[ + { + "model_name": "groq-pinned", + "litellm_params": { + "model": "llama-3.3-70b-versatile", + "api_base": "https://api.groq.com/openai/v1", + "api_key": "fake-key", + }, + }, + { + "model_name": "deepseek-pinned", + "litellm_params": { + "model": "deepseek-chat", + "api_base": "https://api.deepseek.com/v1", + "api_key": "fake-key", + }, + }, + ] + ) + + model_list = router.get_model_list() + assert model_list is not None + assert {m["model_name"] for m in model_list} == {"groq-pinned", "deepseek-pinned"} + + def test_bare_model_with_unknown_api_base_still_raises(self): + with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"): + litellm.Router( + model_list=[ + { + "model_name": "mystery", + "litellm_params": { + "model": "some-unknown-model", + "api_base": "https://llm.internal.example.com/v1", + "api_key": "fake-key", + }, + } + ] + ) + + def test_explicit_custom_llm_provider_beats_api_base_endpoint_match(self): + router = litellm.Router( + model_list=[ + { + "model_name": "openai-via-gateway", + "litellm_params": { + "model": "gpt-3.5-turbo", + "custom_llm_provider": "openai", + "api_base": "https://api.groq.com/openai/v1", + "api_key": "fake-key", + }, + } + ] + ) + + deployment = router.get_deployment_by_model_group_name("openai-via-gateway") + assert deployment is not None + assert deployment.litellm_params.custom_llm_provider == "openai" + +# ===================================================================== +# anthropic_messages mid-stream-fallback helpers, added for #24004 +# (mid-stream fallback not supported for anthropic_messages route type). +# +# anthropic_messages goes through _ageneric_api_call_with_fallbacks rather +# than _acompletion, so its returned iterator was never wrapped by the chat +# completions fallback handler: an SSE `event: error` frame from a native +# Anthropic/Bedrock passthrough passed through to the client silently, and a +# MidStreamFallbackError raised by the completion-bridge path's +# CustomStreamWrapper (e.g. a Vertex AI transport drop) propagated +# unhandled. +# +# Targets the helpers introduced on Router: +# - _aanthropic_messages_streaming_iterator +# - _aanthropic_messages_fallback_attempt +# - _aanthropic_messages_with_streaming_fallbacks +# - _dispatch_generic_call_type +# ===================================================================== + + +async def _anthropic_messages_empty_generator(): + return + yield # pragma: no cover - makes this an async generator + + +def _anthropic_messages_make_wrapper() -> FallbackAwareAnthropicMessagesStream: + """A minimal wrapper for tests that call _aanthropic_messages_fallback_attempt + directly, bypassing _aanthropic_messages_streaming_iterator.""" + return FallbackAwareAnthropicMessagesStream(_anthropic_messages_empty_generator(), object()) + + +def _anthropic_messages_make_router() -> Router: + return Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": { + "model": "anthropic/claude-sonnet-4-5", + "api_key": "sk-test", + }, + }, + { + "model_name": "fallback", + "litellm_params": { + "model": "bedrock/anthropic.claude-sonnet-4-5", + }, + }, + ] + ) + + +class _AnthropicMessagesFakeByteStream: + """Minimal AsyncIterator[bytes], carrying _hidden_params like + AnthropicMessagesStreamingResponse does.""" + + def __init__(self, chunks: list) -> None: + self._chunks = list(chunks) + self._hidden_params = {"additional_headers": {"x-amzn-requestid": "req-1"}} + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self) -> bytes: + if not self._chunks: + raise StopAsyncIteration + return self._chunks.pop(0) + + async def aclose(self) -> None: + self.closed = True + + +class _AnthropicMessagesRaisingByteStream: + """Simulates the completion-bridge path: no error SSE chunk is ever + yielded, the underlying CustomStreamWrapper raises MidStreamFallbackError + directly out of the iterator instead (a Vertex AI transport drop).""" + + def __init__(self, chunks: list, error: Exception) -> None: + self._chunks = list(chunks) + self._error = error + self._hidden_params: dict = {} + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self) -> bytes: + if self._chunks: + return self._chunks.pop(0) + raise self._error + + async def aclose(self) -> None: + self.closed = True + + +class _AnthropicMessagesFallbackByteStream: + def __init__(self, chunks: list, hidden_params: dict | None = None) -> None: + self._chunks = list(chunks) + self._hidden_params = hidden_params if hidden_params is not None else {} + + def __aiter__(self): + return self + + async def __anext__(self) -> bytes: + if not self._chunks: + raise StopAsyncIteration + return self._chunks.pop(0) + + +def _anthropic_messages_overloaded_error_chunk() -> bytes: + return ( + b"event: error\n" + b'data: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' + ) + + +def _anthropic_messages_invalid_request_error_chunk() -> bytes: + return ( + b"event: error\n" + b'data: {"type": "error", "error": {"type": "invalid_request_error", "message": "bad request"}}\n\n' + ) + + +def _anthropic_messages_rate_limit_error_chunk() -> bytes: + return ( + b"event: error\n" + b'data: {"type": "error", "error": {"type": "rate_limit_error", "message": "Too many requests"}}\n\n' + ) + + +def _anthropic_messages_content_chunk(text: str = "hi") -> bytes: + payload = f'{{"type": "content_block_delta", "delta": {{"type": "text_delta", "text": "{text}"}}}}' + return f"event: content_block_delta\ndata: {payload}\n\n".encode() + + +def _anthropic_messages_message_start_chunk() -> bytes: + """A lifecycle/bookkeeping frame Anthropic sends before any real content - + routinely the very first event before an overload error.""" + return b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_1"}}\n\n' + + +def _anthropic_messages_ping_chunk() -> bytes: + return b'event: ping\ndata: {"type": "ping"}\n\n' + + +# -------- _aanthropic_messages_streaming_iterator (passthrough) -------- + + +@pytest.mark.asyncio +async def test_anthropic_messages_streaming_iterator_passthrough(): + """Without any error chunk, the wrapper forwards every chunk unchanged + and carries the source iterator's _hidden_params through (so response + headers like Bedrock's request-id keep flowing to the client).""" + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFakeByteStream( + [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] + ) + + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) + + collected = [chunk async for chunk in wrapped] + assert collected == [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] + assert wrapped._hidden_params["additional_headers"]["x-amzn-requestid"] == "req-1" + + +@pytest.mark.asyncio +async def test_anthropic_messages_streaming_iterator_flushes_buffered_lifecycle_frames_in_order(): + """Regression: lifecycle frames held back to guard against a mid-stream + fallback must still reach the client, in order, once real content + arrives - buffering them for the fallback-safety check must not silently + drop them on the happy path.""" + router = _anthropic_messages_make_router() + message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' + source = _AnthropicMessagesFakeByteStream( + [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] + ) + + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) + + collected = [chunk async for chunk in wrapped] + assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] + + +@pytest.mark.asyncio +async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_stream_end(): + """Regression: if the primary stream ends with only lifecycle frames and + no content and no error, the buffered frames must still reach the + client rather than being silently swallowed.""" + router = _anthropic_messages_make_router() + message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' + source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), message_stop]) + + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) + + collected = [chunk async for chunk in wrapped] + assert collected == [_anthropic_messages_message_start_chunk(), message_stop] + + with pytest.raises(StopAsyncIteration): + await wrapped.__anext__() + + +@pytest.mark.asyncio +async def test_anthropic_messages_content_coalesced_with_error_in_one_physical_chunk_skips_fallback(): + """Greptile review round: transport-level buffering can coalesce a real + content_block_delta and a following retriable error into ONE physical + read from the source iterator. Since the whole chunk (content and error + together) is forwarded to the client atomically, the client genuinely + receives the content - so no fallback must be attempted, exactly as if + the two events had arrived as separate reads.""" + router = _anthropic_messages_make_router() + coalesced_chunk = _anthropic_messages_content_chunk("partial") + _anthropic_messages_overloaded_error_chunk() + source = _AnthropicMessagesFakeByteStream([coalesced_chunk]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [coalesced_chunk] + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_anthropic_messages_ping_keepalive_never_buffered_or_forwarded(): + """Bugbot regression: a `ping` keepalive carries no content and must be + dropped outright before any real content arrives, rather than buffered - + otherwise a slow-starting connection sending many pings could grow the + pre-content buffer without bound.""" + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFakeByteStream( + [_anthropic_messages_ping_chunk(), _anthropic_messages_content_chunk("hi")] + ) + + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) + collected = [chunk async for chunk in wrapped] + + assert _anthropic_messages_ping_chunk() not in collected + assert collected == [_anthropic_messages_content_chunk("hi")] + + +@pytest.mark.asyncio +async def test_anthropic_messages_pre_content_buffer_cap_forces_commit(): + """Bugbot regression: a hostile or pathological upstream that never emits + real content or an error must not grow the pre-content lifecycle buffer + without bound - hitting MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS commits + to the primary stream early, exactly as real content arriving would.""" + router = _anthropic_messages_make_router() + lifecycle_chunk = _anthropic_messages_message_start_chunk() + error_chunk = _anthropic_messages_overloaded_error_chunk() + chunks = [lifecycle_chunk] * (MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + 5) + [error_chunk] + source = _AnthropicMessagesFakeByteStream(chunks) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) + collected = [chunk async for chunk in wrapped] + + mock_fallback.assert_not_awaited() + assert collected.count(lifecycle_chunk) == MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + 5 + assert collected[-1] == error_chunk + + +@pytest.mark.asyncio +async def test_anthropic_messages_ping_coalesced_with_content_in_one_physical_chunk_is_forwarded(): + """Greptile/Bugbot regression: transport-level buffering can coalesce a + `ping` keepalive and a real content_block_delta into ONE physical read. + The pre-content ping-drop must only discard PURE ping frames - dropping + the whole coalesced chunk would silently lose generated content.""" + router = _anthropic_messages_make_router() + coalesced_chunk = _anthropic_messages_ping_chunk() + _anthropic_messages_content_chunk("hi") + source = _AnthropicMessagesFakeByteStream([coalesced_chunk]) + + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) + collected = [chunk async for chunk in wrapped] + + assert collected == [coalesced_chunk] + + +@pytest.mark.asyncio +async def test_anthropic_messages_ping_coalesced_with_retriable_error_still_falls_back(): + """Greptile/Bugbot regression: a physical chunk coalescing a `ping` with a + retriable `event: error` must not be discarded as a keepalive - the error + inside it must still trigger the mid-stream fallback.""" + router = _anthropic_messages_make_router() + coalesced_chunk = _anthropic_messages_ping_chunk() + _anthropic_messages_overloaded_error_chunk() + source = _AnthropicMessagesFakeByteStream([coalesced_chunk]) + fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) + collected = [chunk async for chunk in wrapped] + + mock_fallback.assert_awaited_once() + assert collected == [_anthropic_messages_content_chunk("fallback answer")] + + +# -------- _aanthropic_messages_fallback_attempt -------- + + +@pytest.mark.asyncio +async def test_aanthropic_messages_fallback_attempt_yields_fallback_stream(): + """Direct-call regression: the fallback-attempt helper re-enters the + Router's fallback chain and forwards whatever the fallback produces.""" + router = _anthropic_messages_make_router() + fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) + error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ) as mock_fallback: + collected = [ + chunk + async for chunk in router._aanthropic_messages_fallback_attempt( + error, + {"model": "primary", "messages": [{"role": "user", "content": "hi"}]}, + _anthropic_messages_make_wrapper(), + ) + ] + + assert collected == [_anthropic_messages_content_chunk("fallback answer")] + mock_fallback.assert_awaited_once() + assert mock_fallback.await_args.kwargs["e"] is error + + +@pytest.mark.asyncio +async def test_aanthropic_messages_fallback_attempt_raises_original_exception_on_double_failure(): + """Direct-call regression: when the fallback attempt itself fails with a + MidStreamFallbackError wrapping a real provider exception, that real + exception must surface rather than the internal wrapper exception.""" + router = _anthropic_messages_make_router() + error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") + original_exception = litellm.APIError( + status_code=503, message="fallback also overloaded", llm_provider="bedrock", model="fallback" + ) + fallback_failure = MidStreamFallbackError( + message="fallback failed", model="fallback", llm_provider="bedrock", original_exception=original_exception + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(side_effect=fallback_failure), + ): + with pytest.raises(litellm.APIError) as exc_info: + async for _ in router._aanthropic_messages_fallback_attempt( + error, {"model": "primary"}, _anthropic_messages_make_wrapper() + ): + pass + + assert exc_info.value is original_exception + + +@pytest.mark.asyncio +async def test_aanthropic_messages_fallback_attempt_yields_non_streaming_fallback_response(): + """Bugbot regression: a fallback that resolves to a non-streaming + response (no __aiter__, e.g. an agentic tool-use interception loop) must + be synthesized into a valid SSE byte sequence, not yielded as a raw dict + into a byte stream - the generator is typed AsyncGenerator[bytes, None] + and every item reaching the client must be a real SSE frame.""" + router = _anthropic_messages_make_router() + error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") + non_streaming_response = {"id": "msg_1", "type": "message", "content": [{"type": "text", "text": "hi"}]} + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=non_streaming_response), + ): + collected = [ + item + async for item in router._aanthropic_messages_fallback_attempt( + error, {"model": "primary"}, _anthropic_messages_make_wrapper() + ) + ] + + assert all(isinstance(item, bytes) for item in collected) + event_types = [item.split(b"\n")[0].removeprefix(b"event: ") for item in collected] + assert event_types == [ + b"message_start", + b"content_block_start", + b"content_block_delta", + b"content_block_stop", + b"message_delta", + b"message_stop", + ] + assert b'"text": "hi"' in collected[2] + + +@pytest.mark.asyncio +async def test_aanthropic_messages_fallback_attempt_reraises_plain_exception_on_double_failure(): + """Direct-call regression: when the fallback attempt fails with a plain + exception (not a MidStreamFallbackError), that exception itself must + propagate unchanged.""" + router = _anthropic_messages_make_router() + error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") + fallback_failure = ValueError("no healthy deployments") + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(side_effect=fallback_failure), + ): + with pytest.raises(ValueError, match="no healthy deployments") as exc_info: + async for _ in router._aanthropic_messages_fallback_attempt( + error, {"model": "primary"}, _anthropic_messages_make_wrapper() + ): + pass + + assert exc_info.value is fallback_failure + + +# -------- _aanthropic_messages_with_streaming_fallbacks -------- + + +@pytest.mark.asyncio +async def test_aanthropic_messages_with_streaming_fallbacks_non_streaming_passthrough(): + """A non-streaming response (plain dict) is returned unchanged, never wrapped.""" + router = _anthropic_messages_make_router() + plain_response = {"id": "msg_1", "type": "message"} + + async def fake_original(**_kwargs): + return plain_response + + with patch.object( + router, + "_ageneric_api_call_with_fallbacks", + new=AsyncMock(return_value=plain_response), + ): + out = await router._aanthropic_messages_with_streaming_fallbacks( + original_function=fake_original, + model="primary", + stream=False, + ) + assert out is plain_response + + +@pytest.mark.asyncio +async def test_aanthropic_messages_with_streaming_fallbacks_wraps_streaming_iterator(): + """A streaming response is wrapped via _aanthropic_messages_streaming_iterator.""" + router = _anthropic_messages_make_router() + streaming_iter = _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk()]) + wrapped_marker = object() + + async def fake_original(**_kwargs): + return streaming_iter + + with ( + patch.object( + router, + "_ageneric_api_call_with_fallbacks", + new=AsyncMock(return_value=streaming_iter), + ), + patch.object( + router, + "_aanthropic_messages_streaming_iterator", + new=AsyncMock(return_value=wrapped_marker), + ) as mock_wrap, + ): + out = await router._aanthropic_messages_with_streaming_fallbacks( + original_function=fake_original, + model="primary", + stream=True, + ) + assert out is wrapped_marker + mock_wrap.assert_awaited_once() + + +# -------- mid-stream error handling -------- + + +@pytest.mark.asyncio +async def test_anthropic_messages_fallback_on_pre_first_chunk_error_event(): + """Regression for #24004: a retriable SSE `event: error` frame + (overloaded_error/internal_server_error) that arrives before any real + content must trigger the router's fallback chain instead of passing + through to the client silently.""" + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFakeByteStream([_anthropic_messages_overloaded_error_chunk()]) + fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary", "messages": [{"role": "user", "content": "hi"}]}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [_anthropic_messages_content_chunk("fallback answer")] + mock_fallback.assert_awaited_once() + raised = mock_fallback.await_args.kwargs["e"] + assert isinstance(raised, MidStreamFallbackError) + assert raised.status_code == 503 + assert raised.is_pre_first_chunk is True + assert source.closed is True + + +@pytest.mark.asyncio +async def test_anthropic_messages_mid_stream_error_preserves_real_status_code(): + """Bugbot regression: the MidStreamFallbackError raised for a detected SSE + `event: error` frame must carry the error's REAL parsed status code + (via original_exception), not silently default to 503 for every error + type - a rate_limit_error (429) must surface as 429, not 503.""" + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFakeByteStream([_anthropic_messages_rate_limit_error_chunk()]) + fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary", "messages": [{"role": "user", "content": "hi"}]}, + ) + [chunk async for chunk in wrapped] + + raised = mock_fallback.await_args.kwargs["e"] + assert isinstance(raised, MidStreamFallbackError) + assert raised.status_code == 429 + assert raised.original_exception is not None + assert raised.original_exception.status_code == 429 + assert raised.original_exception.llm_provider == "anthropic" + + +def test_merge_fallback_hidden_params_direct_call(): + """Direct-call regression: merge_fallback_hidden_params combines the + fallback's hidden params/headers with whatever was already present, + with the fallback's values winning on key collisions.""" + wrapper = FallbackAwareAnthropicMessagesStream( + _anthropic_messages_empty_generator(), + _AnthropicMessagesFakeByteStream([]), # carries {"additional_headers": {"x-amzn-requestid": "req-1"}} + ) + wrapper.merge_fallback_hidden_params( + {"model_id": "fallback-deployment"}, + {"x-amzn-requestid": "req-2", "x-fallback-only": "yes"}, + ) + assert wrapper._hidden_params["model_id"] == "fallback-deployment" + assert wrapper._hidden_params["additional_headers"] == { + "x-amzn-requestid": "req-2", + "x-fallback-only": "yes", + } + + +def test_anthropic_stream_should_drop_pre_content_ping_direct_call(): + ping = _anthropic_messages_ping_chunk() + content = _anthropic_messages_content_chunk("hi") + assert _anthropic_stream_should_drop_pre_content_ping(ping, has_generated_content=False) is True + assert _anthropic_stream_should_drop_pre_content_ping(ping, has_generated_content=True) is False + assert _anthropic_stream_should_drop_pre_content_ping(content, has_generated_content=False) is False + + +def test_is_retriable_anthropic_status_direct_call(): + assert _is_retriable_anthropic_status(429) is True + assert _is_retriable_anthropic_status(503) is True + assert _is_retriable_anthropic_status(500) is True + assert _is_retriable_anthropic_status(400) is False + assert _is_retriable_anthropic_status(404) is False + + +def test_anthropic_stream_should_decline_fallback_direct_call(): + pre_first_chunk_error = MidStreamFallbackError( + message="overloaded", model="primary", llm_provider="anthropic", is_pre_first_chunk=True + ) + post_first_chunk_error = MidStreamFallbackError( + message="overloaded", model="primary", llm_provider="anthropic", is_pre_first_chunk=False + ) + assert _anthropic_stream_should_decline_fallback(False, pre_first_chunk_error) is False + assert _anthropic_stream_should_decline_fallback(True, pre_first_chunk_error) is True + assert _anthropic_stream_should_decline_fallback(False, post_first_chunk_error) is True + + +def test_anthropic_stream_commits_now_direct_call(): + content = _anthropic_messages_content_chunk("hi") + lifecycle_chunk = _anthropic_messages_message_start_chunk() + assert _anthropic_stream_commits_now(content, has_generated_content=False, buffered_chunk_count=0) is True + assert _anthropic_stream_commits_now(content, has_generated_content=True, buffered_chunk_count=0) is False + assert ( + _anthropic_stream_commits_now( + lifecycle_chunk, + has_generated_content=False, + buffered_chunk_count=MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS, + ) + is True + ) + assert ( + _anthropic_stream_commits_now( + lifecycle_chunk, + has_generated_content=False, + buffered_chunk_count=MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS - 1, + ) + is False + ) + + +@pytest.mark.asyncio +async def test_anthropic_messages_fallback_merges_fallback_hidden_params(): + """Bugbot regression: after a successful mid-stream fallback, the + wrapper's _hidden_params must reflect the FALLBACK deployment's own + provider headers (e.g. a different Bedrock request-id), not stay + frozen on the primary's - raw bytes can't carry per-item _hidden_params + the way a ModelResponseStream/ResponsesAPI event can, so the wrapper + itself is the only place left to expose them.""" + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFakeByteStream( + [_anthropic_messages_overloaded_error_chunk()] + ) # carries x-amzn-requestid: req-1 + fallback_stream = _AnthropicMessagesFallbackByteStream( + [_anthropic_messages_content_chunk("fallback answer")], + hidden_params={"additional_headers": {"x-amzn-requestid": "req-2", "x-fallback-only": "yes"}}, + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ): + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + _ = [chunk async for chunk in wrapped] + + headers = wrapped._hidden_params["additional_headers"] + assert headers["x-amzn-requestid"] == "req-2" + assert headers["x-fallback-only"] == "yes" + + +@pytest.mark.asyncio +async def test_aanthropic_messages_with_streaming_fallbacks_deep_copies_nested_metadata(): + """Bugbot regression: a shallow .copy() of kwargs still shares the + nested litellm_metadata/metadata dict objects with the primary attempt. + _update_kwargs_with_deployment mutates that dict in place with + deployment-specific fields, which must not leak into the fallback + request's metadata.""" + router = _anthropic_messages_make_router() + primary_metadata = {"model_group": "primary"} + streaming_iter_kwargs = {} + + async def fake_original(**_kwargs): + # Simulate _update_kwargs_with_deployment mutating the primary's + # litellm_metadata in place, as the real helper does. + primary_metadata["deployment"] = "primary-deployment-object" + return _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk("hi")]) + + with patch.object( + router, + "_aanthropic_messages_streaming_iterator", + new=AsyncMock(side_effect=lambda **kwargs: streaming_iter_kwargs.update(kwargs) or "wrapped"), + ): + with patch.object( + router, + "_ageneric_api_call_with_fallbacks", + new=AsyncMock(side_effect=fake_original), + ): + await router._aanthropic_messages_with_streaming_fallbacks( + original_function=fake_original, + model="primary", + stream=True, + litellm_metadata=primary_metadata, + ) + + fallback_kwargs = streaming_iter_kwargs["initial_kwargs"] + assert fallback_kwargs["litellm_metadata"] is not primary_metadata + assert "deployment" not in fallback_kwargs["litellm_metadata"] + + +@pytest.mark.asyncio +async def test_aanthropic_messages_with_streaming_fallbacks_deep_copies_metadata_field(): + """Same regression as above for the (separate) `metadata` kwarg some + call sites use instead of `litellm_metadata`.""" + router = _anthropic_messages_make_router() + primary_metadata = {"tag": "primary"} + streaming_iter_kwargs = {} + + async def fake_original(**_kwargs): + primary_metadata["deployment"] = "primary-deployment-object" + return _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk("hi")]) + + with patch.object( + router, + "_aanthropic_messages_streaming_iterator", + new=AsyncMock(side_effect=lambda **kwargs: streaming_iter_kwargs.update(kwargs) or "wrapped"), + ): + with patch.object( + router, + "_ageneric_api_call_with_fallbacks", + new=AsyncMock(side_effect=fake_original), + ): + await router._aanthropic_messages_with_streaming_fallbacks( + original_function=fake_original, + model="primary", + stream=True, + metadata=primary_metadata, + ) + + fallback_kwargs = streaming_iter_kwargs["initial_kwargs"] + assert fallback_kwargs["metadata"] is not primary_metadata + assert "deployment" not in fallback_kwargs["metadata"] + + +@pytest.mark.asyncio +async def test_anthropic_messages_fallback_triggers_after_lifecycle_only_frame(): + """Regression: Anthropic routinely sends a message_start lifecycle frame + before an overload error even fires. A lifecycle-only frame (no real + content) must not disqualify the fallback attempt, and must not reach + the client either - forwarding it and then appending the fallback's own + message_start would produce two overlapping message lifecycles on one + SSE stream. The primary's buffered lifecycle frame is discarded and the + client sees only the fallback's own, single, clean lifecycle.""" + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFakeByteStream( + [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] + ) + fallback_message_start = b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_2"}}\n\n' + fallback_stream = _AnthropicMessagesFallbackByteStream( + [fallback_message_start, _anthropic_messages_content_chunk("fallback answer")] + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [fallback_message_start, _anthropic_messages_content_chunk("fallback answer")] + assert collected.count(_anthropic_messages_message_start_chunk()) == 0, ( + "the primary's message_start must never reach the client" + ) + assert sum(1 for c in collected if c.startswith(b"event: message_start")) == 1, ( + "exactly one message_start must reach the client" + ) + mock_fallback.assert_awaited_once() + raised = mock_fallback.await_args.kwargs["e"] + assert raised.is_pre_first_chunk is True + + +@pytest.mark.asyncio +async def test_anthropic_messages_raised_error_after_real_content_does_not_restart_stream(): + """Regression: a MidStreamFallbackError raised directly by the source + iterator (the completion-bridge path's CustomStreamWrapper, e.g. a + transport drop) must not trigger a fallback once real content already + reached the client - that would append a second, overlapping message + lifecycle onto the same SSE stream. The original exception must + propagate to the caller instead.""" + router = _anthropic_messages_make_router() + content = _anthropic_messages_content_chunk("partial answer") + original_exception = litellm.APIError( + status_code=503, + message="stream reset", + llm_provider="vertex_ai", + model="primary", + ) + raised_error = MidStreamFallbackError( + message="stream reset", + model="primary", + llm_provider="vertex_ai", + original_exception=original_exception, + is_pre_first_chunk=False, + ) + source = _AnthropicMessagesRaisingByteStream([content], raised_error) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [] + + async def _consume(): + async for chunk in wrapped: + collected.append(chunk) + + with pytest.raises(litellm.APIError) as exc_info: + await _consume() + + assert collected == [content] + assert exc_info.value is original_exception + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_anthropic_messages_fallback_also_catches_raised_midstream_error(): + """Regression for the completion-bridge path (deployments with no native + /v1/messages endpoint): its CustomStreamWrapper raises + MidStreamFallbackError directly (e.g. on a Vertex AI transport drop) + instead of yielding an SSE error chunk - the wrapper must catch that too.""" + router = _anthropic_messages_make_router() + raised_error = MidStreamFallbackError( + message="stream reset", + model="primary", + llm_provider="vertex_ai", + is_pre_first_chunk=True, + ) + source = _AnthropicMessagesRaisingByteStream([], raised_error) + fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [_anthropic_messages_content_chunk("fallback answer")] + mock_fallback.assert_awaited_once() + assert mock_fallback.await_args.kwargs["e"] is raised_error + + +@pytest.mark.asyncio +async def test_anthropic_messages_non_retriable_client_error_skips_fallback(): + """A 4xx (non-429) error type (e.g. invalid_request_error) is a client + error a fallback attempt cannot fix, so it must be forwarded to the + client as-is rather than burning a fallback attempt.""" + router = _anthropic_messages_make_router() + error_chunk = _anthropic_messages_invalid_request_error_chunk() + source = _AnthropicMessagesFakeByteStream([error_chunk]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [error_chunk] + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_anthropic_messages_post_first_chunk_error_skips_fallback(): + """Once content has already reached the caller, retrying would start a + second, overlapping Anthropic message lifecycle on the same SSE stream - + the error must be forwarded instead of triggering an invisible retry.""" + router = _anthropic_messages_make_router() + content = _anthropic_messages_content_chunk("partial answer") + error_chunk = _anthropic_messages_overloaded_error_chunk() + source = _AnthropicMessagesFakeByteStream([content, error_chunk]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [content, error_chunk] + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_anthropic_messages_non_retriable_error_flushes_buffered_lifecycle_frames(): + """A non-retriable error arriving while lifecycle frames are still + buffered (no content seen yet) must flush those buffered frames before + forwarding the error, so the client still sees the whole primary + attempt rather than losing the buffered message_start silently.""" + router = _anthropic_messages_make_router() + error_chunk = _anthropic_messages_invalid_request_error_chunk() + source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), error_chunk]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [_anthropic_messages_message_start_chunk(), error_chunk] + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_anthropic_messages_raised_error_declined_flushes_buffered_lifecycle_frames(): + """When a raised MidStreamFallbackError is declined (source says content + was not pre-first-chunk) while lifecycle frames are still buffered, they + must be flushed to the client before the exception propagates.""" + router = _anthropic_messages_make_router() + raised_error = MidStreamFallbackError( + message="stream reset", + model="primary", + llm_provider="vertex_ai", + is_pre_first_chunk=False, + ) + source = _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], raised_error) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [] + + async def _consume(): + async for chunk in wrapped: + collected.append(chunk) + + with pytest.raises(MidStreamFallbackError) as exc_info: + await _consume() + + assert collected == [_anthropic_messages_message_start_chunk()] + assert exc_info.value is raised_error + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_anthropic_messages_raised_error_without_original_exception_reraises_itself(): + """When a declined MidStreamFallbackError carries no original_exception, + the bare exception itself must propagate rather than being swallowed.""" + router = _anthropic_messages_make_router() + content = _anthropic_messages_content_chunk("partial answer") + raised_error = MidStreamFallbackError( + message="stream reset", + model="primary", + llm_provider="vertex_ai", + is_pre_first_chunk=False, + ) + source = _AnthropicMessagesRaisingByteStream([content], raised_error) + + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [] + + async def _consume(): + async for chunk in wrapped: + collected.append(chunk) + + with pytest.raises(MidStreamFallbackError) as exc_info: + await _consume() + + assert collected == [content] + assert exc_info.value is raised_error + + +@pytest.mark.asyncio +async def test_anthropic_messages_fallback_also_failing_raises_original_exception(): + """If the fallback attempt itself fails with a MidStreamFallbackError + wrapping a real provider exception, the client must see that real + exception, not the internal MidStreamFallbackError.""" + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFakeByteStream([_anthropic_messages_overloaded_error_chunk()]) + original_exception = litellm.APIError( + status_code=503, + message="fallback also overloaded", + llm_provider="bedrock", + model="fallback", + ) + fallback_failure = MidStreamFallbackError( + message="fallback failed", + model="fallback", + llm_provider="bedrock", + original_exception=original_exception, + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(side_effect=fallback_failure), + ): + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + with pytest.raises(litellm.APIError) as exc_info: + async for _ in wrapped: + pass + + assert exc_info.value is original_exception + + +# -------- _dispatch_generic_call_type -------- + + +@pytest.mark.asyncio +async def test_dispatch_generic_call_type_routes_anthropic_messages_through_streaming_fallbacks(): + router = _anthropic_messages_make_router() + + async def fake_original(**_kwargs): + return {"id": "msg_1"} + + with patch.object( + router, + "_aanthropic_messages_with_streaming_fallbacks", + new=AsyncMock(return_value="anthropic-result"), + ) as mock_anthropic: + out = await router._dispatch_generic_call_type( + call_type="anthropic_messages", + original_function=fake_original, + model="primary", + ) + assert out == "anthropic-result" + mock_anthropic.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_dispatch_generic_call_type_other_call_types_use_generic_fallback(): + router = _anthropic_messages_make_router() + + async def fake_original(**_kwargs): + return {"id": "file_1"} + + with patch.object( + router, + "_ageneric_api_call_with_fallbacks", + new=AsyncMock(return_value="generic-result"), + ) as mock_generic: + out = await router._dispatch_generic_call_type( + call_type="afile_delete", + original_function=fake_original, + model="primary", + ) + assert out == "generic-result" + mock_generic.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_factory_function_anthropic_messages_uses_streaming_fallback_dispatch(): + """anthropic_messages must be wired through the mid-stream-fallback-aware + path rather than the bare generic dispatch every other call type without + special handling uses.""" + router = _anthropic_messages_make_router() + wrapped = router.factory_function(litellm.anthropic_messages, call_type="anthropic_messages") + assert callable(wrapped) + + with patch.object( + router, + "_aanthropic_messages_with_streaming_fallbacks", + new=AsyncMock(return_value="ok"), + ) as mock_anthropic: + result = await wrapped(model="primary") + assert result == "ok" + mock_anthropic.assert_awaited_once() diff --git a/tests/test_litellm/test_together_ai_model_metadata.py b/tests/test_litellm/test_together_ai_model_metadata.py new file mode 100644 index 00000000000..5a0aadf4737 --- /dev/null +++ b/tests/test_litellm/test_together_ai_model_metadata.py @@ -0,0 +1,161 @@ +import json +from pathlib import Path +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + +REPO_ROOT: Final = Path(__file__).parents[2] + +CostMap = dict[str, dict[str, object]] +COST_MAP_ADAPTER: Final = TypeAdapter(CostMap) + +SERVERLESS_CHAT_MODELS: Final = ( + "together_ai/moonshotai/Kimi-K3", + "together_ai/zai-org/GLM-5.2", + "together_ai/deepseek-ai/DeepSeek-V4-Pro", + "together_ai/deepseek-ai/DeepSeek-V4-Pro-0813", + "together_ai/deepseek-ai/DeepSeek-V4-Flash-0731", + "together_ai/moonshotai/Kimi-K2.7-Code", + "together_ai/MiniMaxAI/MiniMax-M3", + "together_ai/thinkingmachines/Inkling", + "together_ai/thinkingmachines/Inkling-Small", + "together_ai/Qwen/Qwen3.8-2.4T-A95B", + "together_ai/Qwen/Qwen3.7-Max", + "together_ai/Qwen/Qwen3.7-Plus", + "together_ai/Qwen/Qwen3.6-Plus", + "together_ai/Qwen/Qwen3.5-9B", + "together_ai/nvidia/nemotron-3-ultra-550b-a55b", + "together_ai/meta-models/Muse-Glimmer-30B", + "together_ai/google/gemma-4-31B-it", + "together_ai/pearl-ai/gemma-4-31b-it", + "together_ai/google/gemma-3n-E4B-it", + "together_ai/arize-ai/qwen-2-1.5b-instruct", + "together_ai/Prism-ML/Ternary-Bonsai-27B", + "together_ai/meta-llama/Llama-Guard-4-12B", + "together_ai/openai/gpt-oss-120b", + "together_ai/openai/gpt-oss-20b", + "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", +) + +DEPRECATED_MODELS: Final = { + "together_ai/Qwen/Qwen3-235B-A22B-Instruct-2507-tput": "2026-07-10", + "together_ai/Qwen/Qwen3.5-397B-A17B": "2026-06-29", + "together_ai/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": "2026-06-04", + "together_ai/moonshotai/Kimi-K2.5": "2026-05-21", + "together_ai/deepseek-ai/DeepSeek-R1": "2026-05-14", + "together_ai/deepseek-ai/DeepSeek-V3.1": "2026-05-14", + "together_ai/Qwen/Qwen3-235B-A22B-Thinking-2507": "2026-04-16", + "together_ai/mistralai/Mixtral-8x7B-Instruct-v0.1": "2026-04-16", + "together_ai/zai-org/GLM-4.5-Air-FP8": "2026-04-02", + "together_ai/zai-org/GLM-4.7": "2026-04-02", + "together_ai/mistralai/Mistral-Small-24B-Instruct-2501": "2026-04-02", + "together_ai/Qwen/Qwen3-Next-80B-A3B-Instruct": "2026-04-02", + "together_ai/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": "2026-03-31", + "together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": "2026-03-06", + "together_ai/moonshotai/Kimi-K2-Instruct-0905": "2026-03-06", + "together_ai/meta-llama/Llama-3.2-3B-Instruct-Turbo": "2026-03-06", + "together_ai/Qwen/Qwen3-Next-80B-A3B-Thinking": "2026-02-25", + "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": "2026-02-25", + "together_ai/Qwen/Qwen3-235B-A22B-fp8-tput": "2026-02-06", + "together_ai/meta-llama/Llama-4-Scout-17B-16E-Instruct": "2026-02-06", + "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo": "2026-02-06", + "together_ai/meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": "2026-02-06", + "together_ai/deepseek-ai/DeepSeek-R1-0528-tput": "2026-02-03", + "together_ai/mistralai/Mistral-7B-Instruct-v0.1": "2025-11-13", + "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo-Free": "2025-11-13", +} + + +@pytest.fixture(scope="module") +def cost_map() -> CostMap: + with open(REPO_ROOT / "model_prices_and_context_window.json") as f: + return COST_MAP_ADAPTER.validate_python(json.load(f)) + + +@pytest.mark.parametrize("model", SERVERLESS_CHAT_MODELS) +def test_together_serverless_chat_model_is_mapped(cost_map: CostMap, model: str): + info = cost_map.get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" + assert info["litellm_provider"] == "together_ai" + assert info["mode"] == "chat" + assert info["input_cost_per_token"] >= 0 + assert info["output_cost_per_token"] >= info["input_cost_per_token"] + assert "deprecation_date" not in info + + routed_model, provider, _, _ = get_llm_provider(model=model) + assert routed_model == model.removeprefix("together_ai/") + assert provider == "together_ai" + + +def test_together_kimi_k3_pricing_and_capabilities(cost_map: CostMap): + info = cost_map["together_ai/moonshotai/Kimi-K3"] + assert info["input_cost_per_token"] == 3e-06 + assert info["output_cost_per_token"] == 1.5e-05 + assert info["max_input_tokens"] == 1048576 + assert info["supports_function_calling"] is True + assert info["supports_tool_choice"] is True + assert info["supports_response_schema"] is True + assert info["supports_vision"] is True + assert info["supports_reasoning"] is True + + +def test_together_glm_52_pricing(cost_map: CostMap): + info = cost_map["together_ai/zai-org/GLM-5.2"] + assert info["input_cost_per_token"] == 1.4e-06 + assert info["output_cost_per_token"] == 4.4e-06 + assert info["supports_function_calling"] is True + assert info["supports_reasoning"] is True + + +def test_together_multilingual_e5_embedding_entry(cost_map: CostMap): + info = cost_map["together_ai/intfloat/multilingual-e5-large-instruct"] + assert info["mode"] == "embedding" + assert info["input_cost_per_token"] == 2e-08 + assert info["max_input_tokens"] == 514 + assert info["output_vector_size"] == 1024 + + +def test_together_llama_33_70b_repriced_to_current_together_rate(cost_map: CostMap): + info = cost_map["together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo"] + assert info["input_cost_per_token"] == 1.04e-06 + assert info["output_cost_per_token"] == 1.04e-06 + assert info["max_input_tokens"] == 131072 + + +@pytest.mark.parametrize("model", sorted(DEPRECATED_MODELS)) +def test_together_deprecated_model_carries_deprecation_date(cost_map: CostMap, model: str): + info = cost_map.get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" + assert info.get("deprecation_date") == DEPRECATED_MODELS[model] + + +def _successor(info: dict[str, object]) -> str | None: + metadata = info.get("metadata") + if not isinstance(metadata, dict): + return None + successor = metadata.get("successor") + return successor if isinstance(successor, str) else None + + +def test_together_successor_metadata_points_at_live_models(cost_map: CostMap): + successors = { + model: successor + for model, info in cost_map.items() + if model.startswith("together_ai/") and (successor := _successor(info)) is not None + } + assert len(successors) >= 10 + for model, successor in successors.items(): + target = cost_map.get(successor) + assert target is not None, f"{model} names successor {successor} that is not in the map" + assert "deprecation_date" not in target, f"{model} names deprecated successor {successor}" + + +def test_together_backup_cost_map_in_sync(cost_map: CostMap): + with open(REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json") as f: + backup = COST_MAP_ADAPTER.validate_python(json.load(f)) + together_main = {k: v for k, v in cost_map.items() if k.startswith("together_ai/")} + together_backup = {k: v for k, v in backup.items() if k.startswith("together_ai/")} + assert together_backup == together_main diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 627811a7f1d..7165bb53812 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22805 + "limit": 22804 }, "LIT002": { - "limit": 26873 + "limit": 26872 }, "LIT003": { "limit": 269 @@ -30,7 +30,7 @@ "limit": 16673 }, "LIT011": { - "limit": 5588 + "limit": 5587 }, "LIT012": { "limit": 4510 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx index fd8ed22867b..c64b5d7cb5c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx @@ -6,13 +6,15 @@ import { toast } from "@/lib/toast"; import { Card, CardContent } from "@/components/ui/card"; import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; +import { GuardrailMode } from "@/components/guardrails/types"; +import { formatGuardrailMode } from "./guardrail_info_helpers"; interface GuardrailItem { guardrail_id?: string; guardrail_name: string | null; litellm_params: { guardrail: string; - mode: string; + mode: GuardrailMode; default_on: boolean; }; guardrail_info: Record | null; @@ -171,7 +173,9 @@ const GuardrailTestPlayground: React.FC = ({
Mode: - {guardrail.litellm_params.mode} + + {formatGuardrailMode(guardrail.litellm_params.mode)} +
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx index b4c29bd9c40..7e59abf8e3d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx @@ -18,7 +18,7 @@ import GuardrailTestPlayground from "./GuardrailTestPlayground"; import { toast } from "@/lib/toast"; import { Guardrail } from "@/components/guardrails/types"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; -import { getGuardrailLogoAndName } from "./guardrail_info_helpers"; +import { formatGuardrailMode, getGuardrailLogoAndName } from "./guardrail_info_helpers"; import { CustomCodeModal } from "./custom_code"; import GuardrailGarden from "./guardrail_garden"; import { TeamGuardrailsTab } from "./TeamGuardrailsTab"; @@ -211,7 +211,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole { label: "Name", value: guardrailToDelete?.guardrail_name }, { label: "ID", value: guardrailToDelete?.guardrail_id, code: true }, { label: "Provider", value: providerDisplayName }, - { label: "Mode", value: guardrailToDelete?.litellm_params.mode }, + { label: "Mode", value: formatGuardrailMode(guardrailToDelete?.litellm_params.mode) }, { label: "Default On", value: guardrailToDelete?.litellm_params.default_on ? "Yes" : "No", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx index ec3d05a6907..53f1b1a03d3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx @@ -15,7 +15,7 @@ import { } from "@/components/ui/dropdown-menu"; import { cn } from "@/lib/cva.config"; -import { getGuardrailLogoAndName } from "./guardrail_info_helpers"; +import { formatGuardrailMode, getGuardrailLogoAndName } from "./guardrail_info_helpers"; import { Logo } from "@/components/molecules/logo/Logo"; const CONFIG_DELETE_HINT = "Config guardrails are defined in the config file and cannot be deleted from the dashboard."; @@ -117,9 +117,14 @@ export const getGuardrailTableColumns = ({ header: "Mode", size: 130, enableSorting: false, - cell: ({ row }) => ( - {row.original.litellm_params.mode} - ), + cell: ({ row }) => { + const mode = formatGuardrailMode(row.original.litellm_params.mode); + return ( + + {mode || "-"} + + ); + }, }, { id: "default_on", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx index 4f1ac3e7d7a..2b90a1d8cbc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx @@ -82,6 +82,36 @@ describe("Guardrail Info", () => { expect(getByText("Settings")).toBeInTheDocument(); }); + it("should render a tag-based mode object rather than crashing the detail view", async () => { + vi.mocked(networking.getGuardrailInfo).mockResolvedValue({ + guardrail_id: "123", + guardrail_name: "Test Guardrail", + litellm_params: { + guardrail: "bedrock", + mode: { tags: { "Service-Type: internal-service": "post_call" }, default: ["pre_call", "post_call"] }, + default_on: true, + }, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + guardrail_definition_location: "database", + }); + + vi.mocked(networking.getGuardrailUISettings).mockResolvedValue({ + supported_entities: [], + supported_actions: [], + pii_entity_categories: [], + supported_modes: ["pre_call", "post_call"], + }); + + vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); + + const { findAllByText } = render( + {}} accessToken="123" isAdmin={true} />, + ); + + expect(await findAllByText("pre_call, post_call (tag-based)")).not.toHaveLength(0); + }); + it("should render the provider logo from the bundled guardrail logo map", async () => { vi.mocked(networking.getGuardrailInfo).mockResolvedValue({ guardrail_id: "123", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx index f1b5a0c61ec..226b7b6064d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx @@ -33,6 +33,7 @@ import { import ContentFilterManager, { formatContentFilterDataForAPI } from "./content_filter/ContentFilterManager"; import CustomCodeModal, { EditGuardrailData } from "./custom_code/CustomCodeModal"; import { + formatGuardrailMode, getGuardrailLogoAndName, guardrail_provider_map, skipSystemMessageToChoice, @@ -559,7 +560,9 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose,

Mode

-

{guardrailData.litellm_params?.mode || "-"}

+

+ {formatGuardrailMode(guardrailData.litellm_params?.mode) || "-"} +

{guardrailData.litellm_params?.default_on ? "Default On" : "Default Off"} @@ -856,7 +859,7 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose,

Mode

-
{guardrailData.litellm_params?.mode || "-"}
+
{formatGuardrailMode(guardrailData.litellm_params?.mode) || "-"}

Default On

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx index ec910673b8f..c5e07fe9624 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx @@ -14,6 +14,7 @@ import { choiceToSkipSystemForCreate, skipToolMessageToChoice, choiceToSkipToolForCreate, + formatGuardrailMode, } from "./guardrail_info_helpers"; describe("guardrail_info_helpers", () => { @@ -210,6 +211,34 @@ describe("guardrail_info_helpers", () => { }); }); + describe("formatGuardrailMode", () => { + it("renders a single mode and a list of modes", () => { + expect(formatGuardrailMode("pre_call")).toBe("pre_call"); + expect(formatGuardrailMode(["pre_call", "post_call"])).toBe("pre_call, post_call"); + }); + + it("flattens a tag-based mode object into deduped modes instead of returning it verbatim", () => { + const mode = { + tags: { "Service-Type: internal-service": "post_call", "Service-Type: batch": ["during_call", "post_call"] }, + default: ["pre_call", "post_call"], + }; + + expect(formatGuardrailMode(mode)).toBe("pre_call, post_call, during_call (tag-based)"); + }); + + it("handles a tag-based mode with no default and with no tags", () => { + expect(formatGuardrailMode({ tags: { "team: a": "post_call" } })).toBe("post_call (tag-based)"); + expect(formatGuardrailMode({ default: "pre_call" })).toBe("pre_call (tag-based)"); + }); + + it("returns an empty string for missing or unusable modes", () => { + expect(formatGuardrailMode(undefined)).toBe(""); + expect(formatGuardrailMode(null)).toBe(""); + expect(formatGuardrailMode({})).toBe(""); + expect(formatGuardrailMode({ tags: {}, default: null })).toBe(""); + }); + }); + describe("skipSystemMessageToChoice / choiceToSkipSystemForCreate", () => { it("maps API values to form choices and back for create", () => { expect(skipSystemMessageToChoice(undefined)).toBe("inherit"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index 12aaba0d696..83038b8e0e7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -110,6 +110,17 @@ export const toModeArray = (raw: unknown): string[] => { return []; }; +export const formatGuardrailMode = (raw: unknown): string => { + const flat: string[] = toModeArray(raw); + if (flat.length > 0) return flat.join(", "); + if (raw === null || typeof raw !== "object") return ""; + + const { tags, default: fallback } = raw as { tags?: Record; default?: unknown }; + const tagged: string[] = tags && typeof tags === "object" ? Object.values(tags).flatMap(toModeArray) : []; + const modes: string[] = Array.from(new Set([...toModeArray(fallback), ...tagged])); + return modes.length > 0 ? `${modes.join(", ")} (tag-based)` : ""; +}; + // Resolves the supported modes for the selected provider, falling back to the global list export const getSupportedModesForProvider = ( settings: { supported_modes?: string[]; supported_modes_by_provider?: Record } | null, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx index ee619dc7468..561a89a191a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx @@ -46,6 +46,18 @@ describe("GuardrailTable", () => { expect(screen.getByText("m")).toBeInTheDocument(); }); + it("renders a tag-based mode object instead of crashing the table", () => { + const guardrail = makeGuardrail({ + litellm_params: { + guardrail: "bedrock", + mode: { tags: { "Service-Type: internal-service": "post_call" }, default: ["pre_call", "post_call"] }, + default_on: true, + }, + }); + render(); + expect(screen.getByText("pre_call, post_call (tag-based)")).toBeInTheDocument(); + }); + it("deletes a DB guardrail through the actions menu", async () => { const user = userEvent.setup(); const onDeleteClick = vi.fn(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx index 438caa2f5e6..5aec78ba926 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx @@ -381,6 +381,44 @@ describe("MCPServerEdit (true passthrough warning)", () => { }); }); +describe("MCPServerEdit (OAuth authorize temp payload)", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("forwards issuer/authorization_url/token_url/registration_url to the temp OAuth session payload", async () => { + // Without these fields the ephemeral server the temp OAuth session endpoint builds has no + // admin-configured OAuth endpoints on it, discovery falls back to (and fails against) the + // plain server url, and Authorize & Fetch Token 400s with "authorization url is not + // configured" even though the saved server (and the visible form) has all four fields filled in. + render( + , + ); + + await waitFor(() => { + expect(mockOauth.getTemporaryPayload).toBeTruthy(); + }); + const payload = mockOauth.getTemporaryPayload!(); + expect(payload).toBeTruthy(); + expect(payload?.issuer).toBe("https://github.com/login/oauth"); + expect(payload?.authorization_url).toBe("https://github.com/login/oauth/authorize"); + expect(payload?.token_url).toBe("https://github.com/login/oauth/access_token"); + expect(payload?.registration_url).toBe("https://github.com/login/oauth/register"); + }); +}); + describe("MCPServerEdit (auth type switch)", () => { beforeEach(() => { vi.clearAllMocks(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx index 5e79b20825c..8793c45371a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx @@ -282,6 +282,10 @@ const MCPServerEdit: React.FC = ({ credentials: isClientForwardedTokenMode(values.auth_type) ? preservedAdminCredentials(values.credentials) : values.credentials, + issuer: values.issuer, + authorization_url: values.authorization_url, + token_url: values.token_url, + registration_url: values.registration_url, mcp_access_groups: values.mcp_access_groups || mcpServer.mcp_access_groups, static_headers: staticHeaders, command: values.command, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx index 6cb84e643d3..c4a45d88d1f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx @@ -12,6 +12,7 @@ import { } from "@/components/ui/dialog"; import { Separator } from "@/components/ui/separator"; import { CheckCircle2, Info } from "lucide-react"; +import { formatGuardrailMode } from "@/app/(dashboard)/guardrails/_components/guardrail_info_helpers"; interface GuardrailInfo { guardrail_name: string; @@ -163,7 +164,9 @@ const GuardrailSelectionModal: React.FC = ({ {/* Show guardrail type and mode */}
{guardrail.definition?.litellm_params?.guardrail || "unknown"} - {guardrail.definition?.litellm_params?.mode || "unknown"} + + {formatGuardrailMode(guardrail.definition?.litellm_params?.mode) || "unknown"} + {guardrail.definition?.litellm_params?.patterns && ( {guardrail.definition.litellm_params.patterns.length} pattern(s) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/scope_validation.ts b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/scope_validation.ts index 7c49117088c..53a76dc5a1c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/scope_validation.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/scope_validation.ts @@ -1,4 +1,4 @@ -// Mirrors request-time matching (RouteChecks._route_matches_wildcard_pattern): only a +// Mirrors request-time matching (RouteChecks.route_matches_wildcard_pattern): only a // trailing "*" is a wildcard (prefix match). Anything else - including a "?" or a // non-trailing "*" - is compared by exact equality when a request is matched, so it is // treated as a concrete alias that must exist. diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 303ef3cdddc..45a967c0537 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -8,7 +8,12 @@ vi.mock( ); const mockModelInfo = [ - { model_group: "gpt-4", mode: "chat", supports_reasoning: true }, + { + model_group: "gpt-4", + mode: "chat", + supports_reasoning: true, + supported_reasoning_efforts: ["medium", "high", "xhigh"], + }, { model_group: "gpt-3.5-turbo", mode: "chat" }, { model_group: "claude-3-opus", mode: "chat", supports_reasoning: true }, { model_group: "text-embedding-3-small", mode: "embedding" }, @@ -956,3 +961,54 @@ describe("ComplexityRouterConfig reasoning effort gating", () => { ).toHaveTextContent("low"); }); }); + +describe("ComplexityRouterConfig per-model effort filtering", () => { + it("offers only the efforts the model group supports", async () => { + renderWithProviders(); + const user = userEvent.setup(); + await user.click(screen.getByRole("combobox", { name: "Reasoning effort for gpt-4 in the Complex tier" })); + const options = (await screen.findAllByRole("option")).map((option) => option.textContent); + expect(options).toEqual(["Default", "medium", "high", "xhigh"]); + }); + + it("falls back to every effort when the group only reports supports_reasoning", async () => { + renderWithProviders(); + const user = userEvent.setup(); + await user.click( + screen.getByRole("combobox", { name: "Reasoning effort for claude-3-opus in the Reasoning tier" }), + ); + const options = (await screen.findAllByRole("option")).map((option) => option.textContent); + expect(options).toEqual(["Default", "none", "minimal", "low", "medium", "high", "xhigh"]); + }); + + // An empty list is the group's own answer that its deployments share no level, which is different + // from the field being absent, so the control is dropped rather than falling back to every level. + it("offers no effort at all when the group intersects to nothing", () => { + renderWithProviders( + model.model_group !== "claude-3-opus"), + { model_group: "claude-3-opus", mode: "chat", supports_reasoning: true, supported_reasoning_efforts: [] }, + ]} + />, + ); + expect( + screen.queryByRole("combobox", { name: "Reasoning effort for claude-3-opus in the Reasoning tier" }), + ).not.toBeInTheDocument(); + }); + + // Hand-authored configs can carry a level outside the supported set (e.g. max); it must render + // and stay clearable rather than being masked as Default. + it("keeps showing a stored effort outside the supported set", () => { + renderWithProviders( + , + ); + expect(screen.getByRole("combobox", { name: "Reasoning effort for gpt-4 in the Complex tier" })).toHaveTextContent( + "max", + ); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index c98dc20d3bc..78e5b572492 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -13,6 +13,7 @@ import { ModelGroup } from "@/components/llm_calls/fetch_models"; import AdaptiveRoutingConfig from "./AdaptiveRoutingConfig"; import ClassificationMethodConfig from "./ClassificationMethodConfig"; import { + REASONING_EFFORT_OPTIONS, ReasoningEffort, TierModelParamsByTier, pruneTierModelParams, @@ -252,11 +253,16 @@ const ComplexityRouterConfig: React.FC = ({ const derivedDefaultModel = resolveComplexityDefaultModel(value.tiers); const defaultModel = resolveComplexityDefaultModel(value.tiers, value.default_model); - // Embedding models can't serve a chat-completion role, so they're excluded here. - const reasoningModels = new Set( - modelInfo.filter((model) => model.supports_reasoning).map((model) => model.model_group), + // An absent list means the proxy does not send the field yet, so every level is offered as before. + // An empty list is the group's own answer that its deployments share no level, and is left empty. + const effortOptionsByModel: Record = Object.fromEntries( + modelInfo.map((model) => [ + model.model_group, + model.supported_reasoning_efforts ?? (model.supports_reasoning ? [...REASONING_EFFORT_OPTIONS] : []), + ]), ); + // Embedding models can't serve a chat-completion role, so they're excluded here. const modelOptions = modelInfo .filter((model) => model.mode !== "embedding") .map((model) => ({ @@ -367,7 +373,7 @@ const ComplexityRouterConfig: React.FC = ({ handleTierModelEffortChange(tier, model, effort)} /> diff --git a/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.test.tsx b/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.test.tsx new file mode 100644 index 00000000000..218a18b805e --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.test.tsx @@ -0,0 +1,78 @@ +import { describe, expect, it } from "vitest"; +import { tierEffortRows } from "./TierModelEffortRows"; + +describe("tierEffortRows", () => { + it("offers the levels the proxy reports for the model", () => { + const rows = tierEffortRows({ + models: ["gpt-5-mini"], + effortOptionsByModel: { "gpt-5-mini": ["low", "medium", "high"] }, + paramsByModel: undefined, + }); + + expect(rows).toEqual([{ model: "gpt-5-mini", effort: undefined, options: ["low", "medium", "high"] }]); + }); + + it("keeps a row whose model reports no level but already has an effort stored, so it can be cleared", () => { + const rows = tierEffortRows({ + models: ["off-map-model"], + effortOptionsByModel: {}, + paramsByModel: { "off-map-model": { reasoning_effort: "high" } }, + }); + + expect(rows).toEqual([{ model: "off-map-model", effort: "high", options: ["high"] }]); + }); + + it("drops a row whose model reports no level and has nothing stored", () => { + const rows = tierEffortRows({ + models: ["plain-chat-model"], + effortOptionsByModel: { "plain-chat-model": [] }, + paramsByModel: { "plain-chat-model": { temperature: 0.5 } }, + }); + + expect(rows).toEqual([]); + }); + + it("lists a stored level the model no longer reports without duplicating the ones it does", () => { + const rows = tierEffortRows({ + models: ["gpt-5-mini"], + effortOptionsByModel: { "gpt-5-mini": ["low", "high"] }, + paramsByModel: { "gpt-5-mini": { reasoning_effort: "xhigh" } }, + }); + + expect(rows).toEqual([{ model: "gpt-5-mini", effort: "xhigh", options: ["low", "high", "xhigh"] }]); + }); + + it("does not repeat a stored level the model already reports", () => { + const rows = tierEffortRows({ + models: ["gpt-5-mini"], + effortOptionsByModel: { "gpt-5-mini": ["low", "high"] }, + paramsByModel: { "gpt-5-mini": { reasoning_effort: "high" } }, + }); + + expect(rows).toEqual([{ model: "gpt-5-mini", effort: "high", options: ["low", "high"] }]); + }); + + it.each([ + ["an unset key", {}], + ["an explicit null", { reasoning_effort: null }], + ["an empty string", { reasoning_effort: "" }], + ])("reads %s as no stored effort", (_label, params) => { + const rows = tierEffortRows({ + models: ["gpt-5-mini"], + effortOptionsByModel: { "gpt-5-mini": ["low"] }, + paramsByModel: { "gpt-5-mini": params }, + }); + + expect(rows[0].effort).toBeUndefined(); + }); + + it("renders a non-string stored value as a string so the select can show and clear it", () => { + const rows = tierEffortRows({ + models: ["gpt-5-mini"], + effortOptionsByModel: { "gpt-5-mini": ["low"] }, + paramsByModel: { "gpt-5-mini": { reasoning_effort: 3 } }, + }); + + expect(rows).toEqual([{ model: "gpt-5-mini", effort: "3", options: ["low", "3"] }]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.tsx b/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.tsx index 67f583bd894..ec9705b9451 100644 --- a/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.tsx +++ b/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.tsx @@ -2,35 +2,58 @@ import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@ import { SimpleTooltip } from "@/components/ui/tooltip"; import { Info } from "lucide-react"; import React from "react"; -import { REASONING_EFFORT_OPTIONS, ReasoningEffort, TierModelParams } from "./complexity_router_tiers"; +import { ReasoningEffort, TierModelParams } from "./complexity_router_tiers"; const PROVIDER_DEFAULT = "__provider_default__"; -const asEffort = (params: TierModelParams | undefined): ReasoningEffort | undefined => { +const storedEffort = (params: TierModelParams | undefined): ReasoningEffort | undefined => { const stored = params?.reasoning_effort; - if (typeof stored !== "string") return undefined; - return REASONING_EFFORT_OPTIONS.find((option) => option === stored); + if (stored === undefined || stored === null || stored === "") return undefined; + return typeof stored === "string" ? stored : String(stored); }; interface TierModelEffortRowsProps { tierLabel: string; models: string[]; - reasoningModels: ReadonlySet; + effortOptionsByModel: Record; paramsByModel: Record | undefined; onEffortChange: (model: string, effort: ReasoningEffort | undefined) => void; } +export interface TierEffortRow { + model: string; + effort: ReasoningEffort | undefined; + options: string[]; +} + +/** + * A stored effort outside the model's supported set (hand-authored, or capabilities changed since + * it was saved) is listed anyway, so the row renders with its value selected and can be cleared. + * Only a model with no supported level and nothing stored drops out. + */ +export const tierEffortRows = ({ + models, + effortOptionsByModel, + paramsByModel, +}: Pick): TierEffortRow[] => + models + .map((model) => { + const effort = storedEffort(paramsByModel?.[model]); + const supported = effortOptionsByModel[model] ?? []; + const listed = effort !== undefined && !supported.includes(effort) ? [...supported, effort] : supported; + return { model, effort, options: Array.from(new Set(listed)) }; + }) + .filter(({ options }) => options.length > 0); + const TierModelEffortRows: React.FC = ({ tierLabel, models, - reasoningModels, + effortOptionsByModel, paramsByModel, onEffortChange, }) => { - const shown = models.filter( - (model) => reasoningModels.has(model) || Object.keys(paramsByModel?.[model] ?? {}).length > 0, - ); - if (shown.length === 0) return null; + const rows = tierEffortRows({ models, effortOptionsByModel, paramsByModel }); + if (rows.length === 0) return null; return (
@@ -41,18 +64,17 @@ const TierModelEffortRows: React.FC = ({
- {shown.map((model) => ( + {rows.map(({ model, effort, options }) => (
{model}