From 32469bcab9b4cd795ddabcadca9473286a8cb90c Mon Sep 17 00:00:00 2001 From: Gopal Kamineni <72813883+Gopal3746@users.noreply.github.com> Date: Thu, 5 Nov 2026 21:32:19 -0800 Subject: [PATCH 01/41] fix(vertex_ai): batch Claude models in Anthropic's native row shape (#45715) Backport of #45715 to lts/1.108.x. Cherry-picked from a7030a5a2cab8b1ff525e1e47f8e7bb741609d50 (main). Adapted to the line: the Gemini row mapper keeps the line's signature, since deployment_drop_params on it comes from #31915, which is not on the line, and the Claude row reads drop_params with the line's normalize_drop_params, since drops_unsupported_params is not in core_helpers here. --- litellm/llms/vertex_ai/batches/handler.py | 34 +- .../llms/vertex_ai/batches/transformation.py | 67 ++- .../llms/vertex_ai/files/transformation.py | 258 +++++++++-- tests/unit/batches/test_batch_utils.py | 30 ++ .../llms/vertex_ai/batches/test_handler.py | 47 +- .../vertex_ai/batches/test_transformation.py | 60 +++ .../test_vertex_ai_files_transformation.py | 433 +++++++++++++++++- 7 files changed, 860 insertions(+), 69 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index de5fe39cfc9..a039347e824 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -14,8 +14,10 @@ from litellm.litellm_core_utils.url_utils import ( ) from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, + MaskedHTTPStatusError, get_async_httpx_client, get_httpx_client, + mask_sensitive_info, ) from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM @@ -80,6 +82,15 @@ def _vertex_batch_list_payload(response: _VertexBatchListJsonSource) -> dict[str return response.json() +def _vertex_batch_create_error(error: httpx.HTTPStatusError) -> MaskedHTTPStatusError: + body: Final = mask_sensitive_info(error.response.text) + return MaskedHTTPStatusError( + httpx.HTTPStatusError(f"{error}\n{body}", request=error.request, response=error.response), + message=body, + text=body, + ) + + class VertexAIBatchPrediction(VertexLLM): def __init__(self, gcs_bucket_name: str, *args, **kwargs): super().__init__(*args, **kwargs) @@ -165,11 +176,19 @@ class VertexAIBatchPrediction(VertexLLM): headers=headers, ) - response: Final = sync_handler.post( - url=api_base, - headers=headers, - data=json.dumps(vertex_batch_request), - ) + try: + response: Final = sync_handler.post( + url=api_base, + headers=headers, + data=json.dumps(vertex_batch_request), + ) + except httpx.HTTPStatusError as e: + litellm.verbose_logger.error( + "Vertex AI batch create failed: status=%s, body=%s", + e.response.status_code, + e.response.text[:1000], + ) + raise _vertex_batch_create_error(e) from e payload_view: Final[_VertexBatchPayloadView] = {"payload": response.json()} _json_response: Final = payload_view["payload"] @@ -265,13 +284,12 @@ class VertexAIBatchPrediction(VertexLLM): data=json.dumps(vertex_batch_request), ) except httpx.HTTPStatusError as e: - error_body: Final = e.response.text litellm.verbose_logger.error( "Vertex AI batch create failed: status=%s, body=%s", e.response.status_code, - error_body[:1000], + e.response.text[:1000], ) - raise + raise _vertex_batch_create_error(e) from e payload_view: Final[_VertexBatchPayloadView] = {"payload": response.json()} _json_response: Final = payload_view["payload"] diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index a8d02045740..2460278294c 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -9,14 +9,41 @@ from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.llms.vertex_ai.common_utils import ( VertexAIError, + VertexAIModelRoute, convert_vertex_datetime_to_openai_datetime, + get_vertex_ai_model_route, ) +from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest from litellm.types.llms.vertex_ai import * -from litellm.types.llms.vertex_ai import GenerateContentResponseBody +from litellm.types.llms.vertex_ai import GenerateContentResponseBody, VertexPartnerProvider from litellm.types.utils import LiteLLMBatch, ModelInfo, Usage _NATIVE_VERTEX_RESPONSE: Final = TypeAdapter(GenerateContentResponseBody) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +VERTEX_ANTHROPIC_VERSION: Final = "vertex-2023-10-16" +_ANTHROPIC_PUBLISHER_MODEL_PATH: Final = "publishers/anthropic/models/" + + +def is_vertex_anthropic_batch_model(model: str) -> bool: + """ + Whether a Vertex batch model (bare, `vertex_ai/`-prefixed or a publisher path) is a Claude model, decided the + way Vertex chat routes Claude. Vertex batches Claude only as `publishers/anthropic/models/` with rows in + Anthropic's Messages shape + """ + bare_model: Final = model.rsplit("/", 1)[-1] + return _ANTHROPIC_PUBLISHER_MODEL_PATH in model or ( + get_vertex_ai_model_route(bare_model) is VertexAIModelRoute.PARTNER_MODELS + and VertexPartnerProvider.claude.value in bare_model + ) + + +def as_json_object(value: object) -> Mapping[str, object] | None: + return _JSON_OBJECT.validate_python(value) if isinstance(value, dict) else None + + +def is_vertex_anthropic_batch_output_response(response_body: Mapping[str, object]) -> bool: + return response_body.get("type") == "message" and isinstance(response_body.get("usage"), dict) def _int_field(mapping: Mapping[str, object], key: str) -> int: @@ -43,8 +70,22 @@ def is_vertex_embedding_batch_output_response(response_body: Mapping[str, object return isinstance(response_body.get("embedding"), dict) +def is_vertex_anthropic_batch_output_row(vertex_output_row: Mapping[str, object]) -> bool: + """ + Whether a Vertex batch output row came from a Claude (Anthropic Messages) request. Vertex echoes only the + request's top-level `custom_id` on these rows, with an Anthropic `message` response on success and an empty + response plus a `status` on failure, so a Claude input file read back is never mistaken for output + """ + if not ("custom_id" in vertex_output_row and "processed_time" in vertex_output_row): + return False + response: Final = as_json_object(vertex_output_row.get("response")) + return (response is not None and is_vertex_anthropic_batch_output_response(response)) or bool( + vertex_output_row.get("status") + ) + + def is_native_vertex_batch_output_row(row: Mapping[str, object]) -> bool: - return isinstance(row.get("request"), dict) + return isinstance(row.get("request"), dict) or is_vertex_anthropic_batch_output_row(row) class NativeVertexBatchCostCalculator(Protocol): @@ -72,6 +113,9 @@ def _native_vertex_row_usage( response_body: Mapping[str, object], calculate_usage: Callable[[GenerateContentResponseBody], Usage], ) -> Usage | None: + anthropic_usage: Final = as_json_object(response_body.get("usage")) + if is_vertex_anthropic_batch_output_response(response_body) and anthropic_usage is not None: + return VertexAIAnthropicConfig().calculate_usage(usage_object=anthropic_usage, reasoning_content=None) if "usageMetadata" not in response_body: if not is_vertex_embedding_batch_output_response(response_body): return None @@ -97,21 +141,24 @@ def native_vertex_batch_row_stats( """ Usage and cost of one native Vertex predictions.jsonl row, a `{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}, "modelVersion": ...}}` - generateContent object or a `{"request": ..., "response": {"embedding": {...}, "usageMetadata": {...}}}` - embedding object (an embedding row without `usageMetadata` is billed from its documented `tokenCount`). - `model_name` (the deployment model) prices the row unless it is a wildcard, else its own `modelVersion` - does, else the wildcard name so explicit deployment prices still apply; a row without a response, a - generateContent row without `response.usageMetadata`, and a row whose response fails validation are + generateContent object, a `{"request": ..., "response": {"embedding": {...}, "usageMetadata": {...}}}` + embedding object (an embedding row without `usageMetadata` is billed from its documented `tokenCount`) or a + `{"custom_id": ..., "request": ..., "response": {"type": "message", "model": ..., "usage": {...}}}` Claude object. + `model_name` (the deployment model) prices the row unless it is a wildcard, else its own `modelVersion` (Claude + rows: `model`) does, else the wildcard name so explicit deployment prices still apply; a row without a response, + a generateContent row without `response.usageMetadata`, and a row whose response fails validation are None (failed). """ - response_body: Final = row.get("response") - if not isinstance(response_body, dict): + response_body: Final = as_json_object(row.get("response")) + if response_body is None: return None usage: Final = _native_vertex_row_usage(response_body, calculate_usage) if usage is None: return None total_tokens: Final = usage.total_tokens or (usage.prompt_tokens + usage.completion_tokens) - model_version: Final = response_body.get("modelVersion") + model_version: Final = response_body.get( + "model" if is_vertex_anthropic_batch_output_response(response_body) else "modelVersion" + ) deployment_model: Final = model_name if model_name and "*" not in model_name else None model: Final = deployment_model or (model_version if isinstance(model_version, str) else model_name) if model is None: diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 396acd270fa..b64d81e4630 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -19,6 +19,7 @@ from typing_extensions import ReadOnly, Required import litellm from litellm._uuid import uuid +from litellm.constants import DEFAULT_CHAT_COMPLETION_PARAM_VALUES, RESPONSE_FORMAT_TOOL_NAME from litellm.files.types import FileContentStreamingResult from litellm.files.utils import FilesAPIUtils from litellm.litellm_core_utils.cloud_storage_security import ( @@ -30,11 +31,13 @@ from litellm.litellm_core_utils.cloud_storage_security import ( split_configured_cloud_bucket_name, validate_managed_cloud_file_id, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items, normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_data, extract_file_metadata, ) +from litellm.llms.anthropic.chat.transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY from litellm.llms.base_llm.base_utils import map_developer_role_to_system_role from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body @@ -43,7 +46,14 @@ from litellm.llms.base_llm.files.transformation import ( BaseFileUploadStream, LiteLLMLoggingObj, ) -from litellm.llms.vertex_ai.batches.transformation import vertex_embedding_prompt_token_count +from litellm.llms.vertex_ai.batches.transformation import ( + VERTEX_ANTHROPIC_VERSION, + as_json_object, + is_vertex_anthropic_batch_model, + is_vertex_anthropic_batch_output_response, + is_vertex_anthropic_batch_output_row, + vertex_embedding_prompt_token_count, +) from litellm.llms.vertex_ai.common_utils import ( convert_vertex_datetime_to_openai_datetime, get_vertex_ai_fine_tuned_endpoint_id, @@ -55,6 +65,7 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( transform_openai_input_gemini_embed_content, ) +from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.openai import ( AllMessageValues, @@ -75,6 +86,7 @@ from litellm.types.utils import ( ModelResponse, Usage, ) +from litellm.utils import get_non_default_completion_params, get_optional_params from ..common_utils import VertexAIError from ..vertex_llm_base import VertexBase @@ -83,6 +95,7 @@ _GCP_LABEL_VALUE_MAX_LEN: Final = 63 _CUSTOM_ID_RAW_LABEL_PREFIX: Final = "b32_" _VERTEX_BATCH_KEY_FIELD: Final = "key" _MANAGED_GCS_MODEL_PATH_PATTERN: Final = re.compile(r"publishers/[^/]+/models/([^/?]+)") +_PUBLISHER_MODEL_PATH_PATTERN: Final = re.compile(r"publishers/[^/]+/models/") _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM: Final = ( ("outputDimensionality", "output_dimensionality"), ("taskType", "task_type"), @@ -220,10 +233,14 @@ def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, object]) -> str """ Resolve the OpenAI `custom_id` for a Vertex batch output row. - Embedding rows carry it in the top-level `key` field that Vertex echoes back; - `generateContent` rows have no such field, so it is smuggled through request - labels instead (see `_set_litellm_batch_custom_id_labels`). + Claude rows carry it in their own top-level `custom_id` field and embedding rows in + the top-level `key` field that Vertex echoes back; `generateContent` rows have no + such field, so it is smuggled through request labels instead (see + `_set_litellm_batch_custom_id_labels`). """ + anthropic_custom_id: Final = vertex_output_row.get("custom_id") + if anthropic_custom_id is not None: + return str(anthropic_custom_id) key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) if key is not None: return unquote(str(key)) @@ -284,6 +301,12 @@ def _is_vertex_generate_content_batch_output_row(vertex_output_row: Mapping[str, ) +def _is_vertex_chat_batch_output_row(vertex_output_row: Mapping[str, object]) -> bool: + return _is_vertex_generate_content_batch_output_row(vertex_output_row) or is_vertex_anthropic_batch_output_row( + vertex_output_row + ) + + def _try_parse_vertex_batch_output_row(line: bytes) -> _VertexBatchRow | None: try: row: Final = _parse_vertex_batch_output_row(line.decode("utf-8")) @@ -405,6 +428,37 @@ def _openai_batch_output_row( } +def _vertex_anthropic_batch_output_row_to_openai( + custom_id: str, + response: object, + mock_httpx_response: httpx.Response, +) -> _OpenAIBatchOutputRow: + """ + One OpenAI batch output row from a Vertex Claude row's Anthropic `message` response + """ + anthropic_response: Final = as_json_object(response) + if anthropic_response is None or not is_vertex_anthropic_batch_output_response(anthropic_response): + return _openai_batch_output_row( + custom_id=custom_id, + error_code="vertex_ai_error", + error_message="Claude batch row has no message response", + ) + try: + model_response: Final = VertexAIAnthropicConfig().transform_parsed_response( + completion_response=dict(anthropic_response), + raw_response=mock_httpx_response, + model_response=ModelResponse(), + json_mode=True, + ) + except Exception as e: # noqa: BLE001 # one malformed row must fail alone, not the whole output file + return _openai_batch_output_row( + custom_id=custom_id, + error_code="transformation_error", + error_message=f"Failed to transform response: {e}", + ) + return _openai_batch_output_row(custom_id=custom_id, body=model_response.model_dump()) + + def _split_vertex_batch_key(vertex_output_row: Mapping[str, object]) -> tuple[str, int, int]: """ Resolve `(custom_id, index within that custom_id, group size)` for a Vertex batch @@ -679,21 +733,123 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( ) +_CLAUDE_BATCH_ROW_FIXED_PARAMS: Final = frozenset( + ("model", "messages", "custom_llm_provider", "drop_params", "stream", "stream_options") +) +_CLAUDE_BATCH_OPENAI_PARAMS: Final = frozenset(DEFAULT_CHAT_COMPLETION_PARAM_VALUES) - {"metadata"} +_CLAUDE_BATCH_DEPLOYMENT_PARAMS: Final = ( + _CLAUDE_BATCH_OPENAI_PARAMS - _CLAUDE_BATCH_ROW_FIXED_PARAMS - {"extra_headers"} +) | {"drop_params"} + + +def _claude_batch_deployment_params(litellm_params: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType( + {key: value for key, value in litellm_params.items() if key in _CLAUDE_BATCH_DEPLOYMENT_PARAMS} + ) + + +def _names_the_response_format_tool(tool: object) -> bool: + match tool: + case {"function": {"name": str(name)}}: + return name == RESPONSE_FORMAT_TOOL_NAME + case _: + return False + + +def _openai_batch_jsonl_entry_to_vertex_anthropic_row( + custom_id: object, + messages: list[AllMessageValues], + request_params: Mapping[str, object], + model: str, + deployment_params: Mapping[str, object], +) -> Mapping[str, object]: + """ + Transforms an OpenAI chat batch entry into the Anthropic Messages row Vertex batches Claude with, + `{"custom_id": ..., "request": {"messages": [...], "anthropic_version": "vertex-2023-10-16", "max_tokens": ...}}` + """ + if custom_id is None: + raise ValueError("custom_id is required on Claude lines, since Claude output rows are matched back only by it") + config: Final = VertexAIAnthropicConfig() + bare_model: Final = model.rsplit("/", 1)[-1] + params: Final = {**deployment_params, **request_params} + drop_params: Final = litellm.drop_params is True or normalize_drop_params(params.get("drop_params")) is True + chat_params: Final = { + **get_non_default_completion_params(params), + **{key: value for key, value in params.items() if key in _CLAUDE_BATCH_OPENAI_PARAMS}, + } + if any(map(_names_the_response_format_tool, iter_items(chat_params.get("tools") or ()))): + raise ValueError( + f"tool name {RESPONSE_FORMAT_TOOL_NAME!r} is reserved for response_format on Claude, " + "and batch output cannot tell a user tool by that name apart, so rename it" + ) + optional_params: Final[Mapping[str, object]] = get_optional_params( + model=bare_model, + custom_llm_provider="vertex_ai", + drop_params=drop_params, + messages=messages, + **{ # pyright: ignore[reportArgumentType] # batch JSONL bodies are untyped; get_optional_params checks each param as chat does + key: value for key, value in chat_params.items() if key not in _CLAUDE_BATCH_ROW_FIXED_PARAMS + }, + ) + request_state: Final[dict[str, object]] = {"drop_params": drop_params} + anthropic_request: Final[Mapping[str, object]] = config.transform_request( + model=bare_model, + messages=messages, + optional_params={ + **{key: value for key, value in optional_params.items() if key not in ("json_mode", "extra_body")}, + "anthropic_version": VERTEX_ANTHROPIC_VERSION, + }, + litellm_params=request_state, + headers={}, + ) + renamed_tools: Final = as_str_mapping(request_state.get(ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY)) + if renamed_tools: + raise ValueError( + f"tool names {sorted(map(str, renamed_tools.values()))} must match ^[a-zA-Z0-9_-]{{1,128}}$ for Claude on Vertex, " + "and batch output cannot map a renamed tool back, so rename these tools" + ) + extra_body: Final = as_json_object(optional_params.get("extra_body")) + request: Final[Mapping[str, object]] = ( + anthropic_request + if extra_body is None + else { + **anthropic_request, + **config.transform_extra_body( + extra_body=extra_body, request=anthropic_request, model=bare_model, litellm_params={} + ), + } + ) + return {"custom_id": str(custom_id), "request": request} + + def _openai_batch_jsonl_entry_to_vertex_rows( openai_entry: dict[str, Any], map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]], + batch_model: str | None = None, + deployment_params: Mapping[str, object] = MappingProxyType({}), ) -> tuple[Mapping[str, object], ...]: """ Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to. + `batch_model` is the deployment model the batch runs against and wins over the entry's own `body.model`, + which picks the row shape: Anthropic Messages for Claude, `GenerateContentRequest` otherwise. + jsonl body for vertex is {"request": } Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} """ + entry_body: Final = openai_entry.get("body") + model: Final = batch_model or (str(entry_body.get("model") or "") if isinstance(entry_body, dict) else "") + is_anthropic_model: Final = is_vertex_anthropic_batch_model(model) if _is_embeddings_batch_entry(openai_entry): + if is_anthropic_model: + raise VertexAIError( + status_code=400, + message=f"Vertex AI batch embeddings are not supported for Claude models ('{model}')", + ) return _openai_batch_jsonl_entry_to_vertex_embeddings_rows(openai_entry) - if not isinstance(openai_entry.get("body"), dict): + if not isinstance(entry_body, dict): raise TypeError("`body` on batch requests must be a JSON object, but was missing or not an object") openai_request_body: Final = openai_entry["body"] chat_request_body: Final = ( @@ -701,8 +857,21 @@ def _openai_batch_jsonl_entry_to_vertex_rows( if _is_responses_batch_entry(openai_entry) else openai_request_body ) + chat_messages: Final = map_developer_role_to_system_role(chat_request_body.get("messages", [])) + if is_anthropic_model: + return ( + _openai_batch_jsonl_entry_to_vertex_anthropic_row( + custom_id=openai_entry.get("custom_id"), + messages=chat_messages, + request_params={ + key: value for key, value in chat_request_body.items() if key not in ("model", "messages") + }, + model=model, + deployment_params=deployment_params, + ), + ) vertex_request_body: Final = transform_request_body( - messages=map_developer_role_to_system_role(chat_request_body.get("messages", [])), + messages=chat_messages, model=chat_request_body.get("model", ""), optional_params=map_openai_to_vertex_params(chat_request_body), custom_llm_provider="vertex_ai", @@ -854,15 +1023,24 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): self, openai_file_content: FileTypes, map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]], + batch_model: str | None = None, + deployment_params: Mapping[str, object] = MappingProxyType({}), ) -> None: self._openai_file_content = openai_file_content self._map_openai_to_vertex_params = map_openai_to_vertex_params + self._batch_model = batch_model + self._deployment_params = deployment_params def _iter_vertex_jsonl_chunks(self) -> Iterator[bytes]: first = True for lineno, entry in _iter_openai_jsonl_entries(self._openai_file_content): try: - vertex_rows = _openai_batch_jsonl_entry_to_vertex_rows(entry, self._map_openai_to_vertex_params) + vertex_rows = _openai_batch_jsonl_entry_to_vertex_rows( + entry, + self._map_openai_to_vertex_params, + batch_model=self._batch_model, + deployment_params=self._deployment_params, + ) except Exception as e: raise VertexAIError( status_code=400, @@ -918,16 +1096,31 @@ class _RawFileUploadStream(BaseFileUploadStream): return _iter_raw_file_chunks(self._file_content) +def _deployment_batch_model(litellm_params: Mapping[str, object]) -> str | None: + model: Final = litellm_params.get("model") + return model.removeprefix("vertex_ai/") if isinstance(model, str) and model else None + + +def _publisher_model_path(raw_model: str) -> str: + if _PUBLISHER_MODEL_PATH_PATTERN.search(raw_model): + return raw_model + publisher: Final = "anthropic" if is_vertex_anthropic_batch_model(raw_model) else "google" + return f"publishers/{publisher}/models/{raw_model}" + + def _managed_batch_object_name(raw_model: str, *, passthrough: bool) -> str: endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model) - model_path: Final = ( - f"endpoints/{endpoint_id}" - if endpoint_id is not None - else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}") - ) - safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model") + model_path: Final = f"endpoints/{endpoint_id}" if endpoint_id is not None else _publisher_model_path(raw_model) prefix: Final = _PASSTHROUGH_MANAGED_GCS_PREFIX if passthrough else VERTEX_AI_MANAGED_GCS_PREFIX - return f"{prefix}{safe_model_path}/{uuid.uuid4()}" + return f"{prefix}{_safe_managed_model_path(model_path)}/{uuid.uuid4()}" + + +def _safe_managed_model_path(model_path: str) -> str: + if not is_vertex_anthropic_batch_model(model_path): + return sanitize_cloud_object_path(model_path, fallback="model") + return "@".join( + sanitize_cloud_object_path(version_part, fallback="model") for version_part in model_path.split("@") + ) class VertexAIFilesConfig(VertexBase, BaseFilesConfig): @@ -1030,10 +1223,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): fallback_filename="file", ) - def _get_configured_bucket_name(self, litellm_params: dict) -> str: + def _get_configured_bucket_name(self, litellm_params: Mapping[str, object]) -> str: + configured_bucket_name: Final = litellm_params.get("gcs_bucket_name") or litellm_params.get("bucket_name") bucket_name: Final = ( - litellm_params.get("gcs_bucket_name") - or litellm_params.get("bucket_name") + (configured_bucket_name if isinstance(configured_bucket_name, str) else None) or os.getenv("GCS_BATCH_BUCKET_NAME") or os.getenv("GCS_BUCKET_NAME") ) @@ -1047,7 +1240,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): api_key: str | None, model: str, optional_params: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], data: CreateFileRequest, ) -> str: """ @@ -1071,11 +1264,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): raise ValueError("file is required") if purpose is None: raise ValueError("purpose is required") - configured_model: Final = litellm_params.get("model") object_name = self.get_object_name( file_data, purpose, - deployment_model=configured_model if isinstance(configured_model, str) else None, + deployment_model=_deployment_batch_model(litellm_params), passthrough=is_passthrough_batch_upload(data, litellm_params), ) if object_prefix: @@ -1126,7 +1318,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): model: str, create_file_data: CreateFileRequest, optional_params: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> bytes | str | dict: """ 2 Cases: @@ -1157,6 +1349,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): body_stream=_OpenAIToVertexBatchUploadStream( file_data, self._map_openai_to_vertex_params, + batch_model=_deployment_batch_model(litellm_params), + deployment_params=_claude_batch_deployment_params(litellm_params), ), content_type="application/json", ) @@ -1207,7 +1401,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException: return VertexAIError(status_code=status_code, message=error_message, headers=headers) - def _parse_gcs_uri(self, file_id: str, litellm_params: dict | None = None) -> tuple[str, str]: + def _parse_gcs_uri(self, file_id: str, litellm_params: Mapping[str, object] | None = None) -> tuple[str, str]: """ Validate a managed GCS file_id and return (bucket, url-encoded-object-path). """ @@ -1225,7 +1419,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): self, file_id: str, optional_params: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> tuple[str, dict]: bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params) url: Final = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}" @@ -1255,7 +1449,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): self, file_id: str, optional_params: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> tuple[str, dict]: bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params) url: Final = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}" @@ -1298,7 +1492,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): self, file_content_request, optional_params: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> tuple[str, dict]: file_id: Final = file_content_request.get("file_id", "") bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params) @@ -1395,7 +1589,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): headers=MappingProxyType({**headers, "content-length": str(len(transformed_content))}), ) - if not _is_vertex_generate_content_batch_output_row(first_row): + if not _is_vertex_chat_batch_output_row(first_row): return FileContentStreamingResult(stream_iterator=replayed_stream, headers=headers) return FileContentStreamingResult( @@ -1479,10 +1673,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # first line is not valid UTF-8/JSON) raises and falls through to the # passthrough below, leaving the content untouched. first_row: Final = _parse_vertex_batch_output_row(first_line) - if not ( - _is_vertex_embeddings_batch_output_row(first_row) - or _is_vertex_generate_content_batch_output_row(first_row) - ): + if not (_is_vertex_embeddings_batch_output_row(first_row) or _is_vertex_chat_batch_output_row(first_row)): return content context: Final = _new_vertex_batch_output_row_transform_context() @@ -1547,6 +1738,13 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): error_message=status, ) + if is_vertex_anthropic_batch_output_row(vertex_output): + return _vertex_anthropic_batch_output_row_to_openai( + custom_id=custom_id, + response=vertex_output.get("response"), + mock_httpx_response=mock_httpx_response, + ) + # Transform successful response using existing transformation vertex_response: Final = vertex_output.get("response", {}) diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index ed9bceba927..1ee85522ae3 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -2296,6 +2296,36 @@ async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeyp assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")} +@pytest.mark.asyncio +async def test_raw_vertex_claude_rows_are_priced_like_with_the_native_output_flag(monkeypatch): + claude_message: Final = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "Blue."}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 12, "output_tokens": 5}, + } + rows: Final = [ + {"custom_id": "ok", "processed_time": "2026-10-29T03:00:00Z", "response": claude_message, "status": ""}, + {"custom_id": "bad", "processed_time": "2026-10-29T03:00:00Z", "response": {}, "status": "Bad Request"}, + ] + + async def cost_with_flag(flag: bool) -> bu.BatchCostUsageResult: + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", flag, raising=False) + return await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="vertex_ai/claude-sonnet-4-6" + ) + + unflagged: Final = await cost_with_flag(False) + flagged: Final = await cost_with_flag(True) + + assert unflagged.cost > 0 + assert (unflagged.successful_requests, unflagged.failed_requests) == (1, 1) + assert (unflagged.cost, unflagged.usage.total_tokens) == (flagged.cost, flagged.usage.total_tokens) + + @pytest.mark.asyncio async def test_openai_shaped_vertex_rows_keep_the_generic_path_without_flag(monkeypatch): monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) diff --git a/tests/unit/llms/vertex_ai/batches/test_handler.py b/tests/unit/llms/vertex_ai/batches/test_handler.py index 91a2adb9d19..2c246e065a8 100644 --- a/tests/unit/llms/vertex_ai/batches/test_handler.py +++ b/tests/unit/llms/vertex_ai/batches/test_handler.py @@ -30,6 +30,7 @@ from __future__ import annotations import asyncio import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -41,9 +42,11 @@ from litellm.llms.vertex_ai.batches.handler import ( # noqa: E402 VertexAIBatchPrediction, ) from litellm.llms.vertex_ai.common_utils import VertexAIError # noqa: E402 +from litellm.proxy.utils import handle_exception_on_proxy # noqa: E402 from litellm.types.utils import LiteLLMBatch # noqa: E402 HMOD = "litellm.llms.vertex_ai.batches.handler" +_CLIENT_ERROR_429: Final = "Client error '429 Too Many Requests' for url 'https://x/batchPredictionJobs'" TOKEN = "ya29.fake-access-token" PROJECT = "my-project" LOCATION = "us-central1" @@ -370,19 +373,19 @@ def test_create_batch_sync_endpoint_without_deployed_model_raises_400(): client.post.assert_not_called() -def test_create_batch_sync_httpstatuserror_propagates(): - """``HTTPHandler.post`` raises for non-2xx via ``raise_for_status``; the - sync create path must surface that error, not swallow it.""" +def test_create_batch_sync_httpstatuserror_carries_vertex_status_body_and_headers(): h = _make_handler() client = MagicMock() request = httpx.Request("POST", "https://x/batchPredictionJobs") - err_response = httpx.Response(status_code=500, request=request, text="boom") + err_response: Final = httpx.Response( + status_code=429, request=request, text="boom", headers={"retry-after": "30"} + ) client.post.side_effect = httpx.HTTPStatusError( - "boom", request=request, response=err_response + _CLIENT_ERROR_429, request=request, response=err_response ) with patch(f"{HMOD}.get_httpx_client", return_value=client): - with pytest.raises(httpx.HTTPStatusError): + with pytest.raises(httpx.HTTPStatusError) as exc_info: h.create_batch( _is_async=False, create_batch_data=CREATE_DATA, @@ -393,6 +396,11 @@ def test_create_batch_sync_httpstatuserror_propagates(): timeout=600.0, max_retries=None, ) + assert exc_info.value.response.status_code == 429 + assert exc_info.value.response.headers["retry-after"] == "30" + assert handle_exception_on_proxy(exc_info.value).code == "429" + assert (getattr(exc_info.value, "message", None), getattr(exc_info.value, "text", None)) == ("boom", "boom") + assert str(exc_info.value).splitlines() == [_CLIENT_ERROR_429, "boom"] def test_create_batch_input_file_id_without_model_raises_400_before_post(): @@ -800,14 +808,16 @@ def test_cancel_batch_sync_httpstatuserror_logged_and_reraised(): client.get.assert_not_called() -def test_create_batch_async_httpstatuserror_logged_and_reraised(): +def test_create_batch_async_httpstatuserror_carries_vertex_status_body_and_headers(): h = _make_handler() async_client = MagicMock() request = httpx.Request("POST", "https://x/batchPredictionJobs") - err_response = httpx.Response(status_code=500, request=request, text="boom") + err_response: Final = httpx.Response( + status_code=429, request=request, text="boom", headers={"retry-after": "30"} + ) async_client.post = AsyncMock( side_effect=httpx.HTTPStatusError( - "boom", request=request, response=err_response + _CLIENT_ERROR_429, request=request, response=err_response ) ) @@ -825,8 +835,13 @@ def test_create_batch_async_httpstatuserror_logged_and_reraised(): timeout=600.0, max_retries=None, ) - with pytest.raises(httpx.HTTPStatusError): + with pytest.raises(httpx.HTTPStatusError) as exc_info: _run(coro) + assert exc_info.value.response.status_code == 429 + assert exc_info.value.response.headers["retry-after"] == "30" + assert handle_exception_on_proxy(exc_info.value).code == "429" + assert (getattr(exc_info.value, "message", None), getattr(exc_info.value, "text", None)) == ("boom", "boom") + assert str(exc_info.value).splitlines() == [_CLIENT_ERROR_429, "boom"] def test_async_retrieve_batch_non_200_raises(): @@ -1043,15 +1058,7 @@ async def test_vertex_list_batches(monkeypatch): @pytest.mark.asyncio -async def test_vertex_async_create_batch_logs_error_body_on_http_error(): - """ - When Vertex AI returns an HTTP error (e.g. 400), _async_create_batch should - re-raise httpx.HTTPStatusError (not swallow it) and log the response body. - - Before the fix the error body was lost because AsyncHTTPHandler.post() - calls raise_for_status() internally, raising before the handler's own - status-code check could log the body. - """ +async def test_vertex_async_create_batch_httpstatuserror_carries_vertex_error_body(): from litellm.llms.vertex_ai.batches.handler import VertexAIBatchPrediction handler = VertexAIBatchPrediction(gcs_bucket_name="test-bucket") @@ -1081,4 +1088,4 @@ async def test_vertex_async_create_batch_logs_error_body_on_http_error(): ) assert exc_info.value.response.status_code == 400 - assert "gemini-2.0-flash" in exc_info.value.response.text + assert "Do not support publisher model gemini-2.0-flash" in str(exc_info.value) diff --git a/tests/unit/llms/vertex_ai/batches/test_transformation.py b/tests/unit/llms/vertex_ai/batches/test_transformation.py index 064845a7266..1440a57a15f 100644 --- a/tests/unit/llms/vertex_ai/batches/test_transformation.py +++ b/tests/unit/llms/vertex_ai/batches/test_transformation.py @@ -21,6 +21,8 @@ import pytest from litellm.llms.vertex_ai.batches.transformation import ( # noqa: E402 VertexAIBatchTransformation, + is_vertex_anthropic_batch_model, + native_vertex_batch_row_stats, ) from litellm.llms.vertex_ai.common_utils import VertexAIError # noqa: E402 from litellm.types.utils import LiteLLMBatch # noqa: E402 @@ -493,3 +495,61 @@ def test_get_gcs_uri_prefix_keeps_passthrough_segment_so_output_lands_beside_inp T._get_gcs_uri_prefix_from_file(PASSTHROUGH_INPUT_FILE) == "gs://litellm-testing-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash" ) + + +CLAUDE_INPUT_FILE = ( + "gs://litellm-testing-bucket/litellm-vertex-files/publishers/anthropic/models/claude-sonnet-4-6/uuid-1" +) + + +def test_claude_batch_job_targets_the_anthropic_publisher_model(): + job: Final = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": CLAUDE_INPUT_FILE}, vertex_project="p", vertex_location="us-east5" + ) + assert job["model"] == "publishers/anthropic/models/claude-sonnet-4-6" + + +@pytest.mark.parametrize( + ("model", "expected"), + [ + ("claude-sonnet-4-6", True), + ("vertex_ai/claude-opus-4-1@20250805", True), + ("publishers/anthropic/models/claude-haiku-4-5", True), + ("gemini-2.5-flash", False), + ("publishers/google/models/gemini-2.5-flash", False), + ], +) +def test_is_vertex_anthropic_batch_model(model: str, expected: bool): + assert is_vertex_anthropic_batch_model(model) is expected + + +def test_claude_output_row_is_costed_from_anthropic_usage_and_its_own_model(): + def cost_calculator(usage, model, custom_llm_provider=None, model_info=None, vertex_location=None) -> tuple[float, float]: + if model != "claude-sonnet-4-6": + return 0.0, 0.0 + return usage.prompt_tokens * 0.001, usage.completion_tokens * 0.002 + + stats: Final = native_vertex_batch_row_stats( + { + "custom_id": "r1", + "request": {"messages": [{"role": "user", "content": "hi"}], "anthropic_version": "vertex-2023-10-16"}, + "response": { + "id": "msg_vrtx_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 12, "output_tokens": 5}, + }, + }, + None, + model_info=None, + calculate_usage=lambda _: pytest.fail("Claude rows must not go through the Gemini usage parser"), + cost_calculator=cost_calculator, + ) + + assert stats is not None, "a successful Claude row must be costed, not counted as failed" + assert (stats.usage.prompt_tokens, stats.usage.completion_tokens, stats.total_tokens) == (12, 5, 17) + assert (stats.prompt_cost, stats.completion_cost) == pytest.approx((0.012, 0.010)) + assert (stats.prompt_cost, stats.completion_cost) == pytest.approx((12 * 0.001, 5 * 0.002)) diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index dde3df6fcc0..c235e438847 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -5,7 +5,8 @@ Includes tests for Vertex AI batch output transformation to OpenAI format. import json import urllib.parse -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from typing import Final from types import MappingProxyType from urllib.parse import parse_qs, urlparse @@ -13,6 +14,8 @@ import httpx import pytest from unittest.mock import MagicMock +import litellm +from litellm.llms.base_llm.files.transformation import BaseFileUploadStream from litellm.llms.vertex_ai.files.transformation import ( VertexAIFilesConfig, _get_litellm_batch_custom_id_from_labels, @@ -165,6 +168,17 @@ class TestBatchObjectNaming: object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "gemini-2.5-flash"}}]) assert object_name.startswith("litellm-vertex-files/publishers/google/models/gemini-2.5-flash/") + @pytest.mark.parametrize( + ("model", "stored_model"), + [ + ("multimodalembedding@001", "publishers/google/models/multimodalembedding_001/"), + ("claude-haiku-4-5@20251001", "publishers/anthropic/models/claude-haiku-4-5@20251001/"), + ], + ) + def test_only_claude_models_keep_the_at_version_in_their_object_name(self, config, model, stored_model): + object_name: Final = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": model}}]) + assert object_name.startswith(f"litellm-vertex-files/{stored_model}") + def test_should_store_fine_tuned_endpoint_under_endpoints_path(self, config): """A numeric endpoint id must not be filed under publishers/google/models/gemini/, which the batch transformation later mangles into a nonexistent publisher model (LIT-6899).""" @@ -1948,3 +1962,420 @@ class TestVertexEmbeddingsBatchOutputTranslation: content = json.dumps(legacy_row).encode("utf-8") assert config._try_transform_vertex_batch_output_to_openai(content) == content + + +CLAUDE_DEPLOYMENT = {"gcs_bucket_name": "my-bucket", "model": "vertex_ai/claude-sonnet-4-6"} + + +def _claude_chat_entry(custom_id: str | None, content: str, **body: object) -> dict[str, object]: + entry: Final[dict[str, object]] = { + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "claude-sonnet-4-6", "messages": [{"role": "user", "content": content}], **body}, + } + return entry if custom_id is None else {"custom_id": custom_id, **entry} + + +def _uploaded_rows( + config: VertexAIFilesConfig, litellm_params: Mapping[str, object], entries: Sequence[Mapping[str, object]] +) -> list: + upload: Final = config.transform_create_file_request( + model="", + create_file_data={ + "file": ("batch.jsonl", "\n".join(json.dumps(entry) for entry in entries).encode(), "application/jsonl"), + "purpose": "batch", + }, + optional_params={}, + litellm_params=dict(litellm_params), + ) + match upload: + case {"streaming_media_upload": {"body_stream": BaseFileUploadStream() as stream}}: + return [json.loads(line) for line in b"".join(stream.iter_bytes()).splitlines()] + case _: + pytest.fail(f"expected a streamed upload, got {upload!r}") + + +def _claude_output_row(custom_id: str, text: str, content: object = None) -> dict[str, object]: + return { + "custom_id": custom_id, + "status": "", + "processed_time": "2026-10-28T20:37:47.634+00:00", + "response": { + "id": "msg_vrtx_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": text}] if content is None else content, + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 5}, + }, + } + + +class TestVertexClaudeBatch: + """Claude on Vertex batches only as `publishers/anthropic/models/` with Anthropic Messages rows + (https://github.com/BerriAI/litellm/issues/45671)""" + + @pytest.mark.parametrize("claude_model", ["claude-sonnet-4-6", "claude-haiku-4-5@20251001"]) + def test_claude_upload_creates_a_batch_job_for_the_anthropic_publisher_model(self, config, claude_model): + from litellm.llms.vertex_ai.batches.transformation import VertexAIBatchTransformation + + url: Final = config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={**CLAUDE_DEPLOYMENT, "model": f"vertex_ai/{claude_model}"}, + data={"file": ("batch.jsonl", json.dumps(_claude_chat_entry("r1", "hi")).encode()), "purpose": "batch"}, + ) + object_name: Final = parse_qs(urlparse(url).query)["name"][0] + + job: Final = VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": f"gs://my-bucket/{object_name}"}, vertex_project="p", vertex_location="us-east5" + ) + + assert job["model"] == f"publishers/anthropic/models/{claude_model}", object_name + + def test_claude_body_model_without_deployment_is_filed_under_anthropic_publisher(self, config): + object_name: Final = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "claude-haiku-4-5"}}]) + assert object_name.startswith("litellm-vertex-files/publishers/anthropic/models/claude-haiku-4-5/") + + def test_publisher_path_model_is_not_nested_under_google(self, config): + object_name: Final = config._get_gcs_object_name_from_batch_jsonl( + [{"body": {}}], deployment_model="vertex_ai/publishers/anthropic/models/claude-sonnet-4-6" + ) + assert object_name.startswith("litellm-vertex-files/publishers/anthropic/models/claude-sonnet-4-6/") + + def test_claude_rows_are_uploaded_in_anthropic_messages_shape(self, config): + rows: Final = _uploaded_rows( + config, + CLAUDE_DEPLOYMENT, + [_claude_chat_entry("request-1", "Hello!", max_tokens=50, temperature=0.2)], + ) + + assert rows == [ + { + "custom_id": "request-1", + "request": { + "messages": [{"role": "user", "content": [{"type": "text", "text": "Hello!"}]}], + "max_tokens": 50, + "temperature": 0.2, + "anthropic_version": "vertex-2023-10-16", + }, + } + ] + + def test_claude_system_message_becomes_the_anthropic_system_field(self, config): + entry: Final = _claude_chat_entry( + "r1", + "unused", + max_tokens=10, + messages=[{"role": "system", "content": "be brief"}, {"role": "user", "content": "Hello!"}], + ) + + (row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry]) + + assert row["request"]["system"] == [{"type": "text", "text": "be brief"}] + assert [message["role"] for message in row["request"]["messages"]] == ["user"] + + def test_claude_row_drops_an_unsupported_param_under_deployment_drop_params(self, config, monkeypatch): + monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False}) + monkeypatch.setattr(litellm, "drop_params", False) + entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, top_p=0.5) + + (row,) = _uploaded_rows( + config, {**CLAUDE_DEPLOYMENT, "model": "vertex_ai/claude-zeta-9", "drop_params": True}, [entry] + ) + + assert row["request"]["max_tokens"] == 10 + assert "top_p" not in row["request"] + + def test_claude_row_drops_a_param_the_deployment_lists_in_additional_drop_params(self, config, monkeypatch): + monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False}) + monkeypatch.setattr(litellm, "drop_params", False) + entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, top_p=0.5) + + (row,) = _uploaded_rows( + config, + {**CLAUDE_DEPLOYMENT, "model": "vertex_ai/claude-zeta-9", "additional_drop_params": ["top_p"]}, + [entry], + ) + + assert row["request"]["max_tokens"] == 10 + assert "top_p" not in row["request"] + + def test_claude_rows_take_deployment_params_as_defaults_the_row_overrides(self, config): + rows: Final = _uploaded_rows( + config, + {**CLAUDE_DEPLOYMENT, "max_tokens": 64, "temperature": 0.3}, + [_claude_chat_entry("r1", "hi"), _claude_chat_entry("r2", "hi", max_tokens=10)], + ) + + assert [(row["request"]["max_tokens"], row["request"]["temperature"]) for row in rows] == [(64, 0.3), (10, 0.3)] + + def test_claude_row_with_an_unsupported_param_is_refused_without_drop_params(self, config, monkeypatch): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False}) + monkeypatch.setattr(litellm, "drop_params", False) + entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, top_p=0.5) + + with pytest.raises(VertexAIError, match=r"line 1 .*top_p=0.5.*drop_params") as exc_info: + _uploaded_rows(config, {**CLAUDE_DEPLOYMENT, "model": "vertex_ai/claude-zeta-9"}, [entry]) + assert exc_info.value.status_code == 400 + + def test_claude_deployment_wins_over_a_gemini_body_model(self, config): + entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, model="gemini-2.5-flash") + + (row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry]) + + assert row["request"]["anthropic_version"] == "vertex-2023-10-16" + assert "contents" not in row["request"] + + def test_gemini_deployment_still_uploads_generate_content_rows(self, config): + (row,) = _uploaded_rows( + config, + {"gcs_bucket_name": "my-bucket", "model": "vertex_ai/gemini-2.5-flash"}, + [_claude_chat_entry("r1", "hi")], + ) + assert row["request"]["contents"] == [{"role": "user", "parts": [{"text": "hi"}]}] + assert "custom_id" not in row + + def test_claude_embeddings_batch_line_is_rejected_with_400(self): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + entry: Final = { + "custom_id": "e1", + "url": "/v1/embeddings", + "body": {"model": "claude-sonnet-4-6", "input": "hi"}, + } + with pytest.raises(VertexAIError) as exc_info: + _openai_batch_jsonl_entry_to_vertex_rows(entry, VertexAIFilesConfig()._map_openai_to_vertex_params) + assert exc_info.value.status_code == 400 + + def test_claude_output_rows_become_openai_batch_rows(self, config): + content: Final = b"\n".join( + json.dumps(row).encode() for row in (_claude_output_row("a", "first"), _claude_output_row("b", "second")) + ) + + rows: Final = [ + json.loads(line) for line in config._try_transform_vertex_batch_output_to_openai(content).splitlines() + ] + + assert [row["custom_id"] for row in rows] == ["a", "b"] + assert [row["error"] for row in rows] == [None, None] + bodies: Final = [row["response"]["body"] for row in rows] + assert [body["choices"][0]["message"]["content"] for body in bodies] == ["first", "second"] + assert [body["choices"][0]["finish_reason"] for body in bodies] == ["stop", "stop"] + assert [(body["usage"]["prompt_tokens"], body["usage"]["completion_tokens"]) for body in bodies] == [ + (12, 5), + (12, 5), + ] + assert {body["model"] for body in bodies} == {"claude-sonnet-4-6"} + + @pytest.mark.parametrize("failed_first", [False, True]) + def test_failed_claude_output_row_is_an_openai_error_row(self, config, failed_first): + failed_row: Final = { + "custom_id": "bad", + "status": 'Bad Request: {"type": "error", "error": {"type": "invalid_request_error", "message": "too long"}}', + "processed_time": "2026-10-28T20:37:48.413+00:00", + "response": {}, + } + rows_in: Final = (failed_row, _claude_output_row("ok", "fine")) + content: Final = b"\n".join(json.dumps(row).encode() for row in (rows_in if failed_first else rows_in[::-1])) + + rows_out: Final = { + row["custom_id"]: row + for row in ( + json.loads(line) for line in config._try_transform_vertex_batch_output_to_openai(content).splitlines() + ) + } + + assert rows_out["ok"]["response"]["body"]["choices"][0]["message"]["content"] == "fine" + assert rows_out["bad"]["response"] is None + assert rows_out["bad"]["error"]["code"] == "vertex_ai_error" + assert "too long" in rows_out["bad"]["error"]["message"] + + def test_claude_row_keeps_provider_only_params_and_leaves_out_litellm_ones(self, config): + (row,) = _uploaded_rows( + config, + CLAUDE_DEPLOYMENT, + [_claude_chat_entry("r1", "hi", max_tokens=10, top_k=5, drop_params=True, metadata={"team": "a"})], + ) + + assert row["request"]["top_k"] == 5 + assert not {"drop_params", "metadata", "extra_body"} & set(row["request"]) + + def test_claude_row_top_k_follows_drop_params_on_a_model_without_sampling_params(self, config, monkeypatch): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False}) + monkeypatch.setattr(litellm, "drop_params", False) + deployment: Final = {**CLAUDE_DEPLOYMENT, "model": "vertex_ai/claude-zeta-9"} + entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, top_k=5) + + with pytest.raises(VertexAIError, match=r"line 1 .*top_k") as exc_info: + _uploaded_rows(config, deployment, [entry]) + (row,) = _uploaded_rows(config, {**deployment, "drop_params": True}, [entry]) + + assert exc_info.value.status_code == 400 + assert "top_k" not in row["request"] + + @pytest.mark.parametrize("param", [{"n": 2}, {"seed": 7}, {"presence_penalty": 0.5}, {"logprobs": True}]) + def test_claude_row_param_claude_does_not_support_follows_drop_params(self, config, monkeypatch, param): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + monkeypatch.setattr(litellm, "drop_params", False) + (name,) = param + entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, **param) + + with pytest.raises(VertexAIError, match=rf"line 1 .*{name}") as exc_info: + _uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry]) + (row,) = _uploaded_rows(config, {**CLAUDE_DEPLOYMENT, "drop_params": True}, [entry]) + + assert exc_info.value.status_code == 400 + assert name not in row["request"] + + def test_claude_row_refuses_a_tool_name_claude_would_rename(self, config): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + def tool(name: str) -> dict[str, object]: + return {"type": "function", "function": {"name": name, "parameters": {"type": "object", "properties": {}}}} + + renamed: Final = _claude_chat_entry("r1", "hi", max_tokens=10, tools=[tool("get.weather")]) + kept: Final = _claude_chat_entry("r2", "hi", max_tokens=10, tools=[tool("get_weather")]) + + with pytest.raises(VertexAIError, match=r"line 1 .*get\.weather") as exc_info: + _uploaded_rows(config, CLAUDE_DEPLOYMENT, [renamed]) + (row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [kept]) + + assert exc_info.value.status_code == 400 + assert [t["name"] for t in row["request"]["tools"]] == ["get_weather"] + + def test_claude_row_refuses_a_user_tool_named_like_the_response_format_tool(self, config): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + tool: Final = {"type": "function", "function": {"name": "json_tool_call", "parameters": {"type": "object"}}} + entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, tools=[tool]) + + with pytest.raises(VertexAIError, match=r"line 1 .*json_tool_call") as exc_info: + _uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry]) + + assert exc_info.value.status_code == 400 + + def test_claude_row_leaves_out_stream(self, config): + entry: Final = _claude_chat_entry( + "r1", "hi", max_tokens=10, stream=True, stream_options={"include_usage": True} + ) + + (row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry]) + + assert "stream" not in row["request"] + assert "stream_options" not in row["request"] + + def test_claude_row_merges_extra_body_into_the_request(self, config): + (row,) = _uploaded_rows( + config, CLAUDE_DEPLOYMENT, [_claude_chat_entry("r1", "hi", max_tokens=10, extra_body={"top_k": 3})] + ) + + assert row["request"]["top_k"] == 3 + assert "extra_body" not in row["request"] + + def test_claude_response_format_row_round_trips_as_json_content(self, config): + schema: Final = {"type": "object", "properties": {"color": {"type": "string"}}, "required": ["color"]} + (row,) = _uploaded_rows( + config, + CLAUDE_DEPLOYMENT, + [ + _claude_chat_entry( + "r1", + "sky color?", + max_tokens=50, + response_format={"type": "json_schema", "json_schema": {"name": "sky", "schema": schema}}, + ) + ], + ) + assert "json_mode" not in row["request"] + (json_tool,) = row["request"]["tools"] + output_row: Final = _claude_output_row( + "r1", + "unused", + content=[{"type": "tool_use", "id": "toolu_1", "name": json_tool["name"], "input": {"color": "blue"}}], + ) + + (out,) = ( + json.loads(line) + for line in config._try_transform_vertex_batch_output_to_openai( + json.dumps(output_row).encode() + ).splitlines() + ) + + message: Final = out["response"]["body"]["choices"][0]["message"] + assert json.loads(message["content"]) == {"color": "blue"} + assert not message.get("tool_calls") + + def test_claude_output_keeps_a_user_tool_call_as_a_tool_call(self, config): + output_row: Final = _claude_output_row( + "r1", + "unused", + content=[{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Paris"}}], + ) + + (out,) = ( + json.loads(line) + for line in config._try_transform_vertex_batch_output_to_openai( + json.dumps(output_row).encode() + ).splitlines() + ) + + (tool_call,) = out["response"]["body"]["choices"][0]["message"]["tool_calls"] + assert tool_call["function"]["name"] == "get_weather" + assert json.loads(tool_call["function"]["arguments"]) == {"city": "Paris"} + + def test_claude_input_file_read_back_is_returned_unchanged(self, config): + (row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [_claude_chat_entry("r1", "hi", max_tokens=10)]) + content: Final = json.dumps(row).encode() + + assert config._try_transform_vertex_batch_output_to_openai(content) == content + + def test_malformed_claude_message_fails_only_its_own_row(self, config): + malformed: Final = _claude_output_row("bad", "unused", content=5) + content: Final = json.dumps(_claude_output_row("ok", "fine")).encode() + b"\n" + json.dumps(malformed).encode() + + ok_row, bad_row = ( + json.loads(line) for line in config._try_transform_vertex_batch_output_to_openai(content).splitlines() + ) + + assert ok_row["response"]["body"]["choices"][0]["message"]["content"] == "fine" + assert bad_row["custom_id"] == "bad" + assert bad_row["response"] is None + assert bad_row["error"]["code"] == "transformation_error" + + def test_claude_line_without_custom_id_is_a_400_at_upload(self, config): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + with pytest.raises(VertexAIError, match=r"line 2 .*custom_id") as exc_info: + _uploaded_rows( + config, CLAUDE_DEPLOYMENT, [_claude_chat_entry("r1", "hi", max_tokens=10), _claude_chat_entry(None, "hi")] + ) + + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_streamed_claude_output_is_transformed(self, config): + async def chunks(): + yield json.dumps(_claude_output_row("a", "first")).encode() + b"\n" + + result: Final = await config.transform_file_content_stream( + stream_iterator=chunks(), + headers={}, + request_url="https://storage.googleapis.com/storage/v1/b/my-bucket/o/x?alt=media", + logging_obj=MagicMock(), + litellm_params={}, + ) + body: Final = b"".join([chunk async for chunk in result.stream_iterator]) + + (row,) = (json.loads(line) for line in body.splitlines()) + assert row["custom_id"] == "a" + assert row["response"]["body"]["choices"][0]["message"]["content"] == "first" From 7183904ce335a6db08a8734a3665c0417849c9f8 Mon Sep 17 00:00:00 2001 From: Siddharth Date: Fri, 6 Nov 2026 04:32:37 -0800 Subject: [PATCH 02/41] fix(bedrock): drop another model's signed thinking when replaying to a non-Claude Converse target (#50073) Backport of #50073 to lts/1.108.x. Cherry-picked from 2a4f8b1d1918c074e0752fe50464514c1626ebd8 (main). The line's Converse message builders predate _thinking_in_replay_order and the drop_params/model_info arguments, so drop_signed is passed at the line's own thinking call sites and keeps_cache_point takes only drop_signed. --- .../prompt_templates/factory.py | 56 ++++- litellm/llms/bedrock/common_utils.py | 18 +- .../llms/openai/responses/transformation.py | 4 +- ...bedrock_converse_redacted_thinking_wire.py | 71 ++++++ ..._responses_minted_reasoning_replay_wire.py | 22 +- ...llm_core_utils_prompt_templates_factory.py | 228 +++++++++++++++++- .../test_converse_transformation_nova_2.py | 20 +- .../test_openai_responses_transformation.py | 12 +- 8 files changed, 392 insertions(+), 39 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 2613e0bbd28..60bc9cadee4 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -4607,6 +4607,16 @@ def get_assistant_message_block_or_continue_message( raise ValueError(f"Unsupported content type: {type(content_block)}") +def _as_thinking_block(item: object) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock | None: + match item: + case {"type": "thinking" | "redacted_thinking"}: + return cast( # cast-ok: the mapping pattern matched a thinking block's type tag + ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, item + ) + case _: + return None + + class BedrockConverseMessagesProcessor: @staticmethod def _initial_message_setup( @@ -4655,11 +4665,15 @@ class BedrockConverseMessagesProcessor: deployment_model: str | None = None, base_model: str | None = None, ) -> list[BedrockMessageBlock]: - from litellm.llms.bedrock.common_utils import bedrock_converse_accepts_redacted_thinking + from litellm.llms.bedrock.common_utils import ( + bedrock_converse_accepts_redacted_thinking, + bedrock_converse_rejects_signed_thinking, + ) contents: list[BedrockMessageBlock] = [] msg_i = 0 to_claude: Final = bedrock_converse_accepts_redacted_thinking(model, base_model) + drop_signed_thinking: Final = bedrock_converse_rejects_signed_thinking(model, base_model) messages = BedrockConverseMessagesProcessor._initial_message_setup( messages, model, llm_provider, user_continue_message @@ -4811,7 +4825,7 @@ class BedrockConverseMessagesProcessor: if thinking_blocks is not None: converted_thinking_blocks = ( BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks, to_claude=to_claude + thinking_blocks, to_claude=to_claude, drop_signed=drop_signed_thinking ) ) assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( @@ -4825,7 +4839,7 @@ class BedrockConverseMessagesProcessor: if isinstance(element, dict): if element["type"] == "thinking" or element["type"] == "redacted_thinking": thinking_block = BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks=[element], to_claude=to_claude + thinking_blocks=[element], to_claude=to_claude, drop_signed=drop_signed_thinking ) assistants_parts = ( BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( @@ -4852,7 +4866,9 @@ class BedrockConverseMessagesProcessor: model=model, deployment_model=deployment_model, ) - if _cache_point_block is not None: + if _cache_point_block is not None and BedrockConverseMessagesProcessor.keeps_cache_point( + element, drop_signed=drop_signed_thinking + ): assistants_parts.append(_cache_point_block) assistant_content.extend(assistants_parts) elif _assistant_content is not None and isinstance(_assistant_content, str): @@ -4900,6 +4916,7 @@ class BedrockConverseMessagesProcessor: thinking_blocks: Sequence[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock], *, to_claude: bool = False, + drop_signed: bool = False, ) -> list[BedrockContentBlock]: return [ BedrockContentBlock( @@ -4909,8 +4926,25 @@ class BedrockConverseMessagesProcessor: ) for thinking_block in thinking_blocks if not (to_claude and is_encrypted_reasoning_block(thinking_block)) + and not (drop_signed and BedrockConverseMessagesProcessor._is_signed(thinking_block)) ] + @staticmethod + def keeps_cache_point(element: object, *, drop_signed: bool) -> bool: + match _as_thinking_block(element): + case None: + return True + case thinking_block: + return not (drop_signed and BedrockConverseMessagesProcessor._is_signed(thinking_block)) + + @staticmethod + def _is_signed(thinking_block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock) -> bool: + match thinking_block: + case {"signature": signature} if signature: + return True + case _: + return False + @staticmethod def _to_reasoning_content_block( thinking_block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, @@ -5067,11 +5101,15 @@ def _bedrock_converse_messages_pt( - Please ensure that function response turn comes immediately after a function call turn - Conversation blocks and tool result blocks cannot be provided in the same turn. Issue: https://github.com/BerriAI/litellm/issues/6053 """ - from litellm.llms.bedrock.common_utils import bedrock_converse_accepts_redacted_thinking + from litellm.llms.bedrock.common_utils import ( + bedrock_converse_accepts_redacted_thinking, + bedrock_converse_rejects_signed_thinking, + ) contents: list[BedrockMessageBlock] = [] msg_i = 0 to_claude: Final = bedrock_converse_accepts_redacted_thinking(model, base_model) + drop_signed_thinking: Final = bedrock_converse_rejects_signed_thinking(model, base_model) messages = BedrockConverseMessagesProcessor._initial_message_setup( messages, model, llm_provider, user_continue_message @@ -5224,7 +5262,7 @@ def _bedrock_converse_messages_pt( if thinking_blocks is not None: converted_thinking_blocks = ( BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks, to_claude=to_claude + thinking_blocks, to_claude=to_claude, drop_signed=drop_signed_thinking ) ) assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( @@ -5239,7 +5277,7 @@ def _bedrock_converse_messages_pt( if element["type"] == "thinking" or element["type"] == "redacted_thinking": thinking_block = ( BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks=[element], to_claude=to_claude + thinking_blocks=[element], to_claude=to_claude, drop_signed=drop_signed_thinking ) ) assistants_parts = ( @@ -5268,7 +5306,9 @@ def _bedrock_converse_messages_pt( model=model, deployment_model=deployment_model, ) - if _cache_point_block is not None: + if _cache_point_block is not None and BedrockConverseMessagesProcessor.keeps_cache_point( + element, drop_signed=drop_signed_thinking + ): assistants_parts.append(_cache_point_block) assistant_content.extend(assistants_parts) elif _assistant_content is not None and isinstance(_assistant_content, str): diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 999dd46b01a..44e33f1c43b 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1183,15 +1183,29 @@ def bedrock_converse_accepts_redacted_thinking(model: str, base_model: str | Non its own. An ARN that hides its model, like an application inference profile, falls back to the deployment's ``base_model`` and counts as non-Claude without one. """ + return _is_claude(_converse_target_model(model, base_model)) + + +def bedrock_converse_rejects_signed_thinking(model: str, base_model: str | None) -> bool: + if base_model is None and "." not in get_bedrock_base_model(model): + return False + return not _is_claude_behind_any_region_prefix(_converse_target_model(model, base_model)) + + +def _converse_target_model(model: str, base_model: str | None) -> str: if base_model is None or "." in get_bedrock_base_model(model): - return _is_claude(model) - return _is_claude(base_model) + return model + return base_model def _is_claude(model: str) -> bool: return get_bedrock_base_model(model).startswith(("anthropic", "claude")) +def _is_claude_behind_any_region_prefix(model: str) -> bool: + return _is_claude(model) or get_bedrock_base_model(model).partition(".")[2].startswith("anthropic.") + + def bedrock_converse_supports_parallel_tool_use_config(model: str) -> bool: return any( (litellm.model_cost.get(candidate) or {}).get("supports_parallel_tool_use_config") is True diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 464dc799c16..811b5df9933 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -63,7 +63,9 @@ _NO_TOOL_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) _MODEL_FAMILIES_REJECTING_TOP_LEVEL_SCHEMA_COMBINATORS: Final = ("gpt-4", "gpt-3.5", "chatgpt-4o", "o1", "o3", "o4") _PROVIDERS_WITH_OPENAI_SCHEMA_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) -_PROVIDERS_REPLAYING_ONLY_THEIR_OWN_REASONING: Final = _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS +_PROVIDERS_REPLAYING_ONLY_THEIR_OWN_REASONING: Final = frozenset( + {LlmProviders.AZURE, LlmProviders.BEDROCK, LlmProviders.BEDROCK_MANTLE, LlmProviders.OPENAI} +) _ROLES_READING_PROMPT_CACHE_BREAKPOINTS: Final = frozenset({"system", "developer", "user"}) _PART_TYPES_READING_PROMPT_CACHE_BREAKPOINTS: Final = frozenset({"input_text", "input_image", "input_file"}) _INPUT_ITEM: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) diff --git a/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py b/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py index 79ec0545c44..1da52b0b6dc 100644 --- a/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py +++ b/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py @@ -414,3 +414,74 @@ def test_converse_redacted_thinking_replays_to_the_anthropic_api_as_claude_nativ {"type": "redacted_thinking", "data": NATIVE_REDACTED}, {"type": "text", "text": FIRST_ANSWER}, ) + + +_SIGNED_THINKING: Final[dict[str, JsonValue]] = { + "type": "thinking", + "thinking": "Checking the forecast.", + "signature": "scripted-signature-" + "s" * 32, +} +_CONVERSE_SIGNED_PART: Final[dict[str, JsonValue]] = { + "reasoningContent": { + "reasoningText": {"text": _SIGNED_THINKING["thinking"], "signature": _SIGNED_THINKING["signature"]} + } +} + + +@pytest.mark.parametrize( + ("deployment", "replayed_thinking"), + [ + pytest.param({"model": f"bedrock/converse/{NOVA_MODEL_ID}"}, (), id="nova"), + pytest.param( + { + "model": f"bedrock/converse/{APPLICATION_PROFILE_ARN}", + "model_info": {"base_model": f"bedrock/{NOVA_MODEL_ID}"}, + }, + (), + id="application-profile-arn-with-nova-base-model", + ), + pytest.param( + {"model": f"bedrock/converse/{APPLICATION_PROFILE_ARN}"}, + (_CONVERSE_SIGNED_PART,), + id="application-profile-arn-without-base-model", + ), + pytest.param( + { + "model": f"bedrock/converse/{APPLICATION_PROFILE_ARN}", + "model_info": {"base_model": f"bedrock/{CLAUDE_MODEL_ID}"}, + }, + (_CONVERSE_SIGNED_PART,), + id="application-profile-arn-with-claude-base-model", + ), + pytest.param(_CLAUDE, (_CONVERSE_SIGNED_PART,), id="claude"), + ], +) +def test_signed_history_keeps_its_thinking_unless_the_converse_target_is_known_not_to_be_claude( + gateway: Gateway, deployment: dict[str, JsonValue], replayed_thinking: tuple[dict[str, JsonValue], ...] +) -> None: + model_info: Final = deployment.get("model_info") + parameters: Final = {key: value for key, value in deployment.items() if key != "model_info"} + with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model_info=model_info if isinstance(model_info, dict) else None, + api_key=TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + **parameters, + ) + + reply: Final = _chat( + gateway, + { + "model": model, + "max_tokens": 256, + "messages": [ + {"role": "user", "content": QUESTION}, + {"role": "assistant", "content": FIRST_ANSWER, "thinking_blocks": [_SIGNED_THINKING]}, + {"role": "user", "content": FOLLOW_UP}, + ], + }, + ) + + assert reply.content == FINAL_ANSWER + assert _replayed_assistant_turn(wire).content == (*replayed_thinking, {"text": FIRST_ANSWER}) diff --git a/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py b/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py index 2253443dfc0..4ff5871ff97 100644 --- a/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py +++ b/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py @@ -77,7 +77,18 @@ _AZURE_AI_OPENAI_HOST: Final = _Deployment( "/openai/v1/responses?api-version=preview", strips_message_status=True, ) -_DROPPING: Final = (_OPENAI, _AZURE, _AZURE_AI_OPENAI_HOST) +_BEDROCK_MANTLE: Final = _Deployment( + "bedrock_mantle", "bedrock_mantle/openai.gpt-oss-120b", "synthetic-mantle-key", "/v1/responses" +) +_BEDROCK: Final = _Deployment( + "bedrock", + "bedrock/openai.gpt-oss-120b-1:0", + "synthetic-bedrock-key", + "/openai/v1/responses", + MappingProxyType({"aws_region_name": "us-east-1"}), + model_info=MappingProxyType({"supported_endpoints": ["/v1/responses"]}), +) +_DROPPING: Final = (_OPENAI, _AZURE, _AZURE_AI_OPENAI_HOST, _BEDROCK_MANTLE, _BEDROCK) _KEEPING: Final = ( _Deployment("litellm_proxy", f"litellm_proxy/{_GPT}", "synthetic-proxy-key", "/responses"), _Deployment("databricks", "databricks/gpt-5.6", "synthetic-databricks-key", "/responses"), @@ -95,15 +106,6 @@ _KEEPING: Final = ( "/v1/responses", types_untyped_items_as_messages=True, ), - _Deployment("bedrock_mantle", "bedrock_mantle/openai.gpt-oss-120b", "synthetic-mantle-key", "/v1/responses"), - _Deployment( - "bedrock", - "bedrock/openai.gpt-oss-120b-1:0", - "synthetic-bedrock-key", - "/openai/v1/responses", - MappingProxyType({"aws_region_name": "us-east-1"}), - model_info=MappingProxyType({"supported_endpoints": ["/v1/responses"]}), - ), *( _Deployment(slug, f"{slug}/{model}", f"synthetic-{slug}-key", "/responses") for slug, model in ( diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index ad216548d08..f0b7f9995c4 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -68,7 +68,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.llms.vertex_ai.gemini.transformation import gemini_convert_messages_with_history from litellm.types.llms.bedrock import BedrockConverseReasoningContentBlock -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage READ_FILE: Final[list[object]] = [{"name": "read_file", "description": "Read the contents of a file."}] @@ -422,6 +422,46 @@ def test_anthropic_messages_pt_drops_the_stream_index_from_replayed_thinking_blo ] +@pytest.mark.parametrize( + ("model", "expected_content"), + ( + ( + "global.anthropic.claude-haiku-5-5", + [ + {"reasoningContent": {"reasoningText": {"text": "pondering", "signature": "signed-reasoning"}}}, + {"text": "hello"}, + ], + ), + ("global.zai.glm-5.3", [{"text": "hello"}]), + ("amazon.nova-pro-v1:0", [{"text": "hello"}]), + ("converse/us.openai.gpt-6-luna", [{"text": "hello"}]), + ), +) +def test_bedrock_converse_replays_signed_thinking_only_to_claude( + model: str, expected_content: Sequence[Mapping[str, object]] +) -> None: + body: Final = litellm.AmazonConverseConfig().transform_request( + model=model, + messages=[ + ChatCompletionUserMessage(role="user", content="hi"), + ChatCompletionAssistantMessage( + role="assistant", + content="hello", + thinking_blocks=[{"type": "thinking", "thinking": "pondering", "signature": "signed-reasoning"}], + ), + ChatCompletionUserMessage(role="user", content="continue"), + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert body["messages"] == [ + {"role": "user", "content": [{"text": "hi"}]}, + {"role": "assistant", "content": expected_content}, + {"role": "user", "content": [{"text": "continue"}]}, + ] + + async def test_anthropic_bedrock_thinking_blocks_with_none_content(): """ Test the specific function that processes thinking_blocks when content is None @@ -695,6 +735,192 @@ async def test_bedrock_converse_keeps_dropping_redacted_thinking_for_non_claude_ assert assistant_content == [{"text": "hello"}] +_HIDDEN_MODEL_ARN: Final = "arn:aws:bedrock:us-west-2:123456789012:application-inference-profile/abc123" + + +async def _converse_messages( + messages: list[dict[str, object]], + model: str, + base_model: str | None, + use_async: bool, + assistant_continue_message: str | None = None, +) -> list[dict[str, object]]: + return ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, + model=model, + llm_provider="bedrock", + base_model=base_model, + assistant_continue_message=assistant_continue_message, + ) + if use_async + else _bedrock_converse_messages_pt( + messages=messages, + model=model, + llm_provider="bedrock", + base_model=base_model, + assistant_continue_message=assistant_continue_message, + ) + ) + + +def _assistant_with_thinking(placement: str, block: dict[str, object]) -> dict[str, object]: + return ( + {"role": "assistant", "content": "hello", "thinking_blocks": [block]} + if placement == "thinking_blocks" + else {"role": "assistant", "content": [block, {"type": "text", "text": "hello"}]} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("placement", ["thinking_blocks", "content"]) +@pytest.mark.parametrize( + "signature", ["claude-signature", encrypted_reasoning_signature("gAAAA_1")], ids=["claude", "encrypted"] +) +@pytest.mark.parametrize( + ("model", "base_model"), + [ + ("converse/us.openai.gpt-6-luna", None), + ("us.amazon.nova-pro-v1:0", None), + ("us.amazon.nova-pro-v1:0", "anthropic.claude-sonnet-4-6"), + (_HIDDEN_MODEL_ARN, "openai.gpt-oss-120b-1:0"), + ], +) +async def test_bedrock_converse_drops_signed_thinking_for_a_known_non_claude_target( + use_async: bool, placement: str, signature: str, model: str, base_model: str | None +): + block: Final = {"type": "thinking", "thinking": "private plan", "signature": signature} + messages: Final = [ + {"role": "user", "content": "hi"}, + _assistant_with_thinking(placement, block), + {"role": "user", "content": "again"}, + ] + + assert await _converse_messages(messages, model, base_model, use_async) == [ + {"role": "user", "content": [{"text": "hi"}]}, + {"role": "assistant", "content": [{"text": "hello"}]}, + {"role": "user", "content": [{"text": "again"}]}, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("placement", ["thinking_blocks", "content"]) +@pytest.mark.parametrize( + ("model", "base_model"), + [ + ("us.anthropic.claude-sonnet-4-6", None), + ("ca.anthropic.claude-haiku-5-5", None), + ("converse/in.anthropic.claude-sonnet-5", None), + ("claude-sonnet-4-6", None), + (_HIDDEN_MODEL_ARN, None), + (_HIDDEN_MODEL_ARN, "anthropic.claude-sonnet-4-6"), + ], +) +async def test_bedrock_converse_keeps_signed_thinking_for_a_target_that_may_be_claude( + use_async: bool, placement: str, model: str, base_model: str | None +): + block: Final = {"type": "thinking", "thinking": "private plan", "signature": "claude-signature"} + messages: Final = [ + {"role": "user", "content": "hi"}, + _assistant_with_thinking(placement, block), + {"role": "user", "content": "again"}, + ] + + result: Final = await _converse_messages(messages, model, base_model, use_async) + + assistant_content: Final = next(m["content"] for m in result if m["role"] == "assistant") + assert assistant_content == [ + {"reasoningContent": {"reasoningText": {"text": "private plan", "signature": "claude-signature"}}}, + {"text": "hello"}, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_bedrock_converse_keeps_the_tool_call_when_dropping_signed_thinking(use_async: bool): + messages: Final = [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": None, + "thinking_blocks": [{"type": "thinking", "thinking": "private plan", "signature": "claude-signature"}], + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ] + + result: Final = await _converse_messages(messages, "converse/us.openai.gpt-6-luna", None, use_async) + + assert [m["role"] for m in result] == ["user", "assistant", "user"] + assert [list(part) for part in result[1]["content"]] == [["toolUse"]] + assert result[1]["content"][0]["toolUse"]["toolUseId"] == "call_1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize( + ("assistant_continue_message", "expected"), + [ + (None, [{"role": "user", "content": [{"text": "hi"}, {"text": "again"}]}]), + ( + "Please continue.", + [ + {"role": "user", "content": [{"text": "hi"}]}, + {"role": "assistant", "content": [{"text": "Please continue."}]}, + {"role": "user", "content": [{"text": "again"}]}, + ], + ), + ], + ids=["merged_into_one_user_turn", "split_by_the_continue_message"], +) +async def test_bedrock_converse_drops_a_turn_that_held_only_signed_thinking( + use_async: bool, assistant_continue_message: str | None, expected: list[dict[str, object]] +): + messages: Final = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "thinking_blocks": [{"type": "thinking", "thinking": "private plan", "signature": "claude-signature"}], + }, + {"role": "user", "content": "again"}, + ] + + result: Final = await _converse_messages( + messages, "converse/us.openai.gpt-6-luna", None, use_async, assistant_continue_message + ) + + assert result == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_bedrock_converse_drops_the_cache_point_of_a_dropped_signed_thinking_block(use_async: bool): + messages: Final = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "private plan", + "signature": "claude-signature", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + {"role": "user", "content": "again"}, + ] + + result: Final = await _converse_messages(messages, "converse/us.amazon.nova-2-lite-v1:0", None, use_async) + + assert result == [{"role": "user", "content": [{"text": "hi"}, {"text": "again"}]}] + + @pytest.mark.parametrize( "thinking_block", [ diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation_nova_2.py b/tests/unit/llms/bedrock/chat/test_converse_transformation_nova_2.py index 58058a2e1d4..f29ad3b5844 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation_nova_2.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation_nova_2.py @@ -637,8 +637,7 @@ class TestNova2MultiTurnMessageTranslation: assert "Simple addition" in texts assert "4." in texts - def test_should_keep_signed_thinking_blocks_as_reasoning_content(self): - """With a signature, reasoning is preserved as a reasoningContent block.""" + def test_should_drop_signed_thinking_blocks_from_another_model(self): bedrock_msgs = self._to_bedrock( [ {"role": "user", "content": "What is 2+2?"}, @@ -653,13 +652,9 @@ class TestNova2MultiTurnMessageTranslation: ] ) assistant = next(m for m in bedrock_msgs if m["role"] == "assistant") - rc_blocks = [b for b in assistant["content"] if "reasoningContent" in b] - assert len(rc_blocks) >= 1 - assert rc_blocks[0]["reasoningContent"]["reasoningText"]["text"] == "math" - assert rc_blocks[0]["reasoningContent"]["reasoningText"]["signature"] == "sig-1" + assert assistant["content"] == [{"text": "4."}] - def test_should_translate_inline_content_list_thinking_type(self): - """content=[{type:'thinking',...},{type:'text',...}] should also round-trip.""" + def test_should_drop_inline_signed_thinking_and_keep_the_answer(self): bedrock_msgs = self._to_bedrock( [ {"role": "user", "content": "Hi"}, @@ -674,11 +669,4 @@ class TestNova2MultiTurnMessageTranslation: ] ) assistant = next(m for m in bedrock_msgs if m["role"] == "assistant") - rc_blocks = [b for b in assistant["content"] if "reasoningContent" in b] - text_blocks = [ - b - for b in assistant["content"] - if "text" in b and "reasoningContent" not in b - ] - assert len(rc_blocks) >= 1 - assert any("Hello!" in b["text"] for b in text_blocks) + assert assistant["content"] == [{"text": "Hello!"}] diff --git a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py index 2fcb70ec40b..0a0f0d33543 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py @@ -13,6 +13,8 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.llms.bedrock.responses.transformation import BedrockOpenAIResponsesConfig +from litellm.llms.bedrock_mantle.responses.transformation import BedrockMantleResponsesAPIConfig from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig from litellm.types.llms.openai import ( @@ -1298,7 +1300,15 @@ class TestOpenAIResponsesAPIConfig: ) return list(bridged.output) - @pytest.mark.parametrize("config", [OpenAIResponsesAPIConfig(), AzureOpenAIResponsesAPIConfig()]) + @pytest.mark.parametrize( + "config", + [ + OpenAIResponsesAPIConfig(), + AzureOpenAIResponsesAPIConfig(), + BedrockOpenAIResponsesConfig(), + BedrockMantleResponsesAPIConfig(), + ], + ) def test_claude_reasoning_minted_by_the_bridge_is_dropped_before_the_history_reaches_openai(self, config): saved_claude_turn = json.loads( json.dumps([item.model_dump() for item in self._claude_turn_bridged_to_responses_output()]) From 534aeb2d88f91904ff9b4e0a4f980d7199bc963b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 12:29:32 -0800 Subject: [PATCH 03/41] fix(bedrock): treat an empty api_base or runtime endpoint as unset Backport to lts/1.108.x ahead of its main PR. Cherry-picked from 9bcb039ac2 (litellm_fix_bedrock_empty_api_base). --- litellm/llms/bedrock/base_aws_llm.py | 29 ++++------------ .../test_bedrock_token_counter.py | 30 +++++++++++++++++ tests/unit/llms/bedrock/test_base_aws_llm.py | 33 +++++++++++++++++++ 3 files changed, 70 insertions(+), 22 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index f1426aa9375..2c72a491b3c 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1526,28 +1526,13 @@ class BaseAWSLLM(SignsRequestsWithAWS): aws_region_name: str, endpoint_type: Literal["runtime", "agent", "agentcore"] | None = "runtime", ) -> tuple[str, str]: - env_aws_bedrock_runtime_endpoint: Final = get_secret("AWS_BEDROCK_RUNTIME_ENDPOINT") - if api_base is not None: - endpoint_url = api_base - elif aws_bedrock_runtime_endpoint is not None and isinstance(aws_bedrock_runtime_endpoint, str): - endpoint_url = aws_bedrock_runtime_endpoint - elif env_aws_bedrock_runtime_endpoint and isinstance(env_aws_bedrock_runtime_endpoint, str): - endpoint_url = env_aws_bedrock_runtime_endpoint - else: - endpoint_url = self._select_default_endpoint_url( - endpoint_type=endpoint_type, - aws_region_name=aws_region_name, - ) - - # Determine proxy_endpoint_url - if aws_bedrock_runtime_endpoint is not None and isinstance(aws_bedrock_runtime_endpoint, str): - proxy_endpoint_url = aws_bedrock_runtime_endpoint - elif env_aws_bedrock_runtime_endpoint and isinstance(env_aws_bedrock_runtime_endpoint, str): - proxy_endpoint_url = env_aws_bedrock_runtime_endpoint - else: - proxy_endpoint_url = endpoint_url - - return endpoint_url, proxy_endpoint_url + runtime_endpoint: Final = aws_bedrock_runtime_endpoint or get_secret_str("AWS_BEDROCK_RUNTIME_ENDPOINT") + endpoint_url: Final = ( + api_base + or runtime_endpoint + or self._select_default_endpoint_url(endpoint_type=endpoint_type, aws_region_name=aws_region_name) + ) + return endpoint_url, runtime_endpoint or endpoint_url def _select_default_endpoint_url( self, diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py index dd0d6429f8a..8749b01ed5f 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -89,6 +89,36 @@ async def test_bedrock_runtime_count_is_kept_when_it_answers() -> None: assert bedrock.posted_hosts() == (RUNTIME_HOST,) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_to_use", "runtime", "expected_hosts", "expected_tokens"), + ( + ("us.anthropic.claude-sonnet-4-6", (200, {"inputTokens": 1353}), (RUNTIME_HOST,), 1353), + ("us.anthropic.claude-opus-5-5", (400, UNSUPPORTED), (RUNTIME_HOST, MANTLE_HOST), 2177), + ), +) +async def test_deployment_with_an_empty_api_base_counts_on_the_default_runtime_host( + model_to_use: str, + runtime: tuple[int, Mapping[str, object]], + expected_hosts: tuple[str, ...], + expected_tokens: int, +) -> None: + bedrock: Final = _Bedrock(runtime=runtime, mantle=(200, {"input_tokens": 2177})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use=model_to_use, + messages=MESSAGES, + contents=None, + deployment={"litellm_params": {**DEPLOYMENT["litellm_params"], "api_base": ""}}, + request_model=model_to_use, + ) + + assert result is not None + assert result.error is False + assert result.total_tokens == expected_tokens + assert bedrock.posted_hosts() == expected_hosts + + @pytest.mark.asyncio @pytest.mark.parametrize( ("model_to_use", "runtime"), diff --git a/tests/unit/llms/bedrock/test_base_aws_llm.py b/tests/unit/llms/bedrock/test_base_aws_llm.py index c755a7e6bc2..d690eda09e8 100644 --- a/tests/unit/llms/bedrock/test_base_aws_llm.py +++ b/tests/unit/llms/bedrock/test_base_aws_llm.py @@ -3848,6 +3848,39 @@ def test_get_runtime_endpoint(base_aws_llm): assert endpoint_url == "https://bedrock-runtime.us-east-1.amazonaws.com" assert proxy_endpoint_url == "https://bedrock-runtime.us-east-1.amazonaws.com" +@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") +@pytest.mark.parametrize( + ("api_base", "aws_bedrock_runtime_endpoint", "env_runtime_endpoint", "expected"), + ( + ("", None, None, "https://bedrock-runtime.us-east-1.amazonaws.com"), + (None, "", None, "https://bedrock-runtime.us-east-1.amazonaws.com"), + ("", "", None, "https://bedrock-runtime.us-east-1.amazonaws.com"), + ("", "https://vpce-runtime.example.com", None, "https://vpce-runtime.example.com"), + ("", None, "https://env-runtime.example.com", "https://env-runtime.example.com"), + (None, "", "https://env-runtime.example.com", "https://env-runtime.example.com"), + ), +) +def test_get_runtime_endpoint_treats_an_empty_setting_as_unset( + base_aws_llm: BaseAWSLLM, + monkeypatch: pytest.MonkeyPatch, + api_base: str | None, + aws_bedrock_runtime_endpoint: str | None, + env_runtime_endpoint: str | None, + expected: str, +) -> None: + if env_runtime_endpoint is None: + monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) + else: + monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", env_runtime_endpoint) + + endpoints: Final = base_aws_llm.get_runtime_endpoint( + api_base=api_base, + aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, + aws_region_name="us-east-1", + ) + + assert endpoints == (expected, expected) + @pytest.fixture def clear_cache(base_aws_llm): """Clear the cache before each test""" From 1e513a4314039f6654c0defbb624b8459b7db772 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 12:56:05 -0800 Subject: [PATCH 04/41] fix(chatgpt): round-trip encrypted reasoning through /v1/messages The chat bridge stored encrypted reasoning on reasoning_items, the Anthropic adapter never read that field, stream chunks that only carried thinking blocks were dropped, and the messages adapter stripped the signed blocks before a chatgpt request could turn them back into reasoning items. Backport of #49076 to lts/1.108.x. Cherry-picked from ea9e29b1a89 (cursor/chatgpt-encrypted-reasoning-messages-1723). The line keeps the reasoning helpers in common_utils and the adapters under pass_through, gains resolve_model_and_provider_for_metadata, and streams redacted_thinking as its own block. Co-authored-by: DDT <141296842+Hugo-DDT@users.noreply.github.com> Co-authored-by: DDT <1786035110@qq.com> --- .../transformation.py | 73 +++++++- .../get_llm_provider_logic.py | 8 + .../prompt_templates/common_utils.py | 82 +++++++++ .../litellm_core_utils/streaming_handler.py | 1 + .../adapters/streaming_iterator.py | 10 +- .../pass_through/adapters/transformation.py | 76 ++++++--- litellm/responses/bridge_check.py | 25 ++- ...responses_transformation_transformation.py | 159 +++++++++++++++++ .../test_streaming_handler.py | 31 ++++ ...al_pass_through_adapters_transformation.py | 160 ++++++++++++++++++ 10 files changed, 591 insertions(+), 34 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 822ef1ea0be..7fe36c2e42a 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -26,9 +26,11 @@ import litellm from litellm import ModelResponse from litellm._logging import verbose_logger from litellm.integrations.anthropic_cache_control_hook import supports_openai_prompt_cache_breakpoint +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( responses_reasoning_items_from_thinking_blocks, + thinking_blocks_from_reasoning_items, with_prompt_cache_breakpoint, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator @@ -45,6 +47,8 @@ from litellm.responses.utils import ResponsesAPIRequestUtils, normalize_response from litellm.types.llms.openai import ( ChatCompletionAnnotation, ChatCompletionReasoningItem, + ChatCompletionRedactedThinkingBlock, + ChatCompletionThinkingBlock, ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, ChatCompletionToolParamFunctionChunk, @@ -54,7 +58,7 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ResponsesAPIStreamEvents, ) -from litellm.types.utils import GenericStreamingChunk, ModelResponseStream +from litellm.types.utils import Delta, GenericStreamingChunk, ModelResponseStream, StreamingChoices if TYPE_CHECKING: from openai.types.responses import ResponseInputImageParam, ResponseOutputItem @@ -70,8 +74,6 @@ if TYPE_CHECKING: AllMessageValues, ChatCompletionFileObject, ChatCompletionImageObject, - ChatCompletionRedactedThinkingBlock, - ChatCompletionThinkingBlock, ChatCompletionToolReferenceObject, OpenAIMessageContentListBlock, ) @@ -443,6 +445,63 @@ def _as_chat_reasoning_items( return cast(list[ChatCompletionReasoningItem], list(reasoning_items)) +def _with_stream_block_index( + block: Mapping[str, object], + index: int, +) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock: + if block.get("type") == "redacted_thinking": + data: Final = block.get("data") + return ChatCompletionRedactedThinkingBlock( + type="redacted_thinking", + data=data if isinstance(data, str) else "", + index=index, + ) + thinking: Final = block.get("thinking") + signature: Final = block.get("signature") + return ChatCompletionThinkingBlock( + type="thinking", + thinking=thinking if isinstance(thinking, str) else "", + signature=signature if isinstance(signature, str) else None, + index=index, + ) + + +def _reasoning_done_thinking_blocks( + item: Mapping[str, object], + output_index: object, +) -> tuple[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, ...]: + index: Final = output_index if isinstance(output_index, int) and not isinstance(output_index, bool) else 0 + return tuple(_with_stream_block_index(block, index) for block in thinking_blocks_from_reasoning_items((item,))) + + +def _reasoning_done_stream_chunk(chunk: Mapping[str, object]) -> ModelResponseStream | None: + event_type: Final = chunk.get("type") + event_name: Final = ( + event_type.value + if isinstance(event_type, ResponsesAPIStreamEvents) + else event_type + if isinstance(event_type, str) + else None + ) + if event_name != ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE: + return None + item: Final = as_str_mapping(chunk.get("item")) + if item is None or item.get("type") != "reasoning": + return None + signed_blocks: Final = _reasoning_done_thinking_blocks(item, chunk.get("output_index")) + if not signed_blocks: + return None + return ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(thinking_blocks=list(signed_blocks)), + finish_reason=None, + ) + ] + ) + + def _shift_annotation(annotation: ChatCompletionAnnotation, offset: int) -> ChatCompletionAnnotation: shifted: Final = { key: value + offset if key in ("start_index", "end_index") and isinstance(value, int) else value @@ -497,6 +556,7 @@ def _fold_choices_into_one( *chain.from_iterable(_message_reasoning_items(m) for m in messages), *(() if trailing_reasoning_item is None else (trailing_reasoning_item,)), ) + encrypted_thinking_blocks: Final = thinking_blocks_from_reasoning_items(reasoning_items) return Choices( message=Message( role=messages[0].role if messages else "assistant", @@ -505,6 +565,7 @@ def _fold_choices_into_one( tool_calls=list(tool_calls) or None, reasoning_content=" ".join(value for value in reasoning_contents if value) or None, reasoning_items=_as_chat_reasoning_items(reasoning_items), + thinking_blocks=list(encrypted_thinking_blocks) or None, provider_specific_fields=_tool_search_provider_fields(tool_search_items), ), finish_reason="tool_calls" if tool_calls else "stop", @@ -1310,10 +1371,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): for summary_block in reasoning_item["summary"] if summary_block.get("text") ) + encrypted_thinking_blocks: Final = thinking_blocks_from_reasoning_items(reasoning_items) message: Final = Message( content=None if tool_search_items else "", reasoning_content=reasoning_content if reasoning_content else None, reasoning_items=_as_chat_reasoning_items(reasoning_items), + thinking_blocks=list(encrypted_thinking_blocks) or None, provider_specific_fields=_tool_search_provider_fields(tool_search_items), ) return Choices(message=message, finish_reason=finish_reason, index=0) @@ -2089,7 +2152,6 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ) ] ) - elif event_type == "response.output_text.delta": # Content part added to output content_part = parsed_chunk.get("delta", None) @@ -2193,9 +2255,12 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): """ verbose_logger.debug("Chat provider: transform_streaming_response called with chunk: %s", chunk) self._remember_served_service_tier(chunk) + reasoning_chunk: Final = _reasoning_done_stream_chunk(chunk) translated: Final = ( self._empty_content_chunk() if self._is_text_delta_after_structured_answer(chunk) + else reasoning_chunk + if reasoning_chunk is not None else OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream( chunk, tool_call_index_map=self._tool_call_index_map ) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index a67c1bb7df8..a9be736360b 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -158,6 +158,14 @@ def declared_authenticating_provider(model: str | None, custom_llm_provider: str return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None +def resolve_model_and_provider_for_metadata(model: str, custom_llm_provider: str | None = None) -> tuple[str, str]: + declared: Final = declared_authenticating_provider(model, custom_llm_provider) + if declared is not None: + return model.removeprefix(f"{declared}/"), declared + resolved_model, resolved_provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) + return resolved_model, resolved_provider + + def inferred_provider(model: str | None) -> str | None: if not model: return None diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index ac5791957b9..790860ca561 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -20,6 +20,7 @@ from pydantic import TypeAdapter import litellm from litellm import verbose_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.router_utils.batch_utils import InMemoryFile from litellm.types.llms.openai import ( AllMessageValues, @@ -2289,6 +2290,87 @@ def encrypted_content_of_block(block: Mapping[str, object]) -> str | None: return encrypted_content_from_signature(_encrypted_reasoning_field(block)) +_REASONING_SUMMARY_SEPARATOR: Final = "\n\n" + + +def _object_sequence(value: object) -> tuple[object, ...]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + return () + return tuple(value) + + +def _iter_mappings(value: object) -> tuple[Mapping[str, object], ...]: + return tuple(mapping for item in _object_sequence(value) if (mapping := as_str_mapping(item)) is not None) + + +def _reasoning_summary_part_text(part: object) -> str: + mapping: Final = as_str_mapping(part) + if mapping is not None: + text: Final = mapping.get("text") + return text if isinstance(text, str) else "" + text_attr: Final = getattr(part, "text", None) + return text_attr if isinstance(text_attr, str) else "" + + +def _reasoning_item_summary_text(item: Mapping[str, object]) -> str: + parts: Final = tuple(_reasoning_summary_part_text(part) for part in _object_sequence(item.get("summary"))) + return _REASONING_SUMMARY_SEPARATOR.join(part for part in parts if part) + + +def thinking_block_from_reasoning_item( + item: object, +) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock | None: + """The chat thinking block that carries one Responses `reasoning` item back to an Anthropic client. + + `encrypted_content` rides the same `litellm_encrypted_reasoning:` signature the Responses + adapter uses. A summary becomes thinking text. Bytes with no summary become `redacted_thinking`, + because an empty thinking string is not something the client can replay. + """ + mapping: Final = as_str_mapping(item) + if mapping is None or mapping.get("type") not in (None, "reasoning"): + return None + encrypted_content: Final = mapping.get("encrypted_content") + if not isinstance(encrypted_content, str) or not encrypted_content: + return None + signature: Final = encrypted_reasoning_signature(encrypted_content) + text: Final = _reasoning_item_summary_text(mapping) + if not text: + return ChatCompletionRedactedThinkingBlock(type="redacted_thinking", data=signature) + return ChatCompletionThinkingBlock(type="thinking", thinking=text, signature=signature) + + +def thinking_blocks_from_reasoning_items( + items: object, +) -> tuple[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, ...]: + return tuple( + block for item in _object_sequence(items) if (block := thinking_block_from_reasoning_item(item)) is not None + ) + + +def _block_encrypted_content(block: object) -> str | None: + mapping: Final = as_str_mapping(block) + if mapping is None: + return None + return encrypted_content_of_block(mapping) + + +def thinking_blocks_with_encrypted_reasoning( + thinking_blocks: object, + reasoning_items: object, +) -> tuple[Mapping[str, object], ...]: + """Thinking blocks already on the message, plus any encrypted reasoning item they do not already carry.""" + direct: Final = _iter_mappings(thinking_blocks) + represented: Final = frozenset( + content for block in direct if (content := _block_encrypted_content(block)) is not None + ) + extra: Final = tuple( + block + for block in thinking_blocks_from_reasoning_items(_iter_mappings(reasoning_items)) + if _block_encrypted_content(block) not in represented + ) + return (*extra, *direct) + + def is_encrypted_reasoning_block(block: object) -> bool: """A thinking or redacted_thinking block carrying Responses API encrypted reasoning. diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 7135b9c1071..d91a67a2d13 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -912,6 +912,7 @@ class CustomStreamWrapper: and model_response.choices[0].delta.role is not None ) or (getattr(model_response.choices[0].delta, "reasoning_items", None) is not None) + or getattr(model_response.choices[0].delta, "thinking_blocks", None) ): return True else: diff --git a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py index 18db0f28338..0f1398ac9b2 100644 --- a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py @@ -24,7 +24,9 @@ from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.exceptions import ContextWindowExceededError, MidStreamFallbackError +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.litellm_core_utils.hidden_params import set_hidden_params +from litellm.litellm_core_utils.prompt_templates.common_utils import thinking_blocks_from_reasoning_items from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -66,6 +68,10 @@ def _optional_attr(obj: object, name: str) -> object: return getattr(obj, name, None) +def _delta_has_encrypted_reasoning(delta: object) -> bool: + return bool(thinking_blocks_from_reasoning_items(attribute_or_none(delta, "reasoning_items"))) + + def _optional_attr_sequence(obj: object, name: str) -> Sequence[object]: value: Final = getattr(obj, name, None) return value if value else () @@ -458,7 +464,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): sent_first_chunk: bool = False sent_content_block_start: bool = False sent_content_block_finish: bool = False - current_content_block_type: Literal["text", "tool_use", "thinking"] = "text" + current_content_block_type: Literal["text", "tool_use", "thinking", "redacted_thinking"] = "text" sent_last_message: bool = False holding_chunk: ContentBlockDelta | None = None holding_stop_reason_chunk: MessageBlockDelta | None = None @@ -1306,6 +1312,8 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): isinstance(b, dict) and not is_empty_unsigned_thinking_block(b) for b in thinking_blocks ): return False + if _delta_has_encrypted_reasoning(delta): + return False return True def _should_start_new_content_block(self, chunk: "ModelResponseStream") -> bool: diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index d1560c5b2ed..62f33bb11fd 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -122,6 +122,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( anthropic_image_source_to_openai_url, parse_tool_call_arguments, reasoning_content_from_thinking_blocks, + thinking_blocks_with_encrypted_reasoning, with_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -219,6 +220,19 @@ def target_supports_mid_conversation_system(model: str | None, custom_llm_provid return supports_mid_conversation_system(model=model, custom_llm_provider=custom_llm_provider) +def _keeps_encrypted_reasoning_for_responses(model: str | None, custom_llm_provider: str | None) -> bool: + from litellm.responses.bridge_check import chat_completion_ends_in_responses_api + + return chat_completion_ends_in_responses_api(model, custom_llm_provider) + + +def _anthropic_thinking_blocks(source: object) -> tuple[Mapping[str, object], ...]: + return thinking_blocks_with_encrypted_reasoning( + getattr(source, "thinking_blocks", None), + getattr(source, "reasoning_items", None), + ) + + def _chat_tool_param(function_chunk: ChatCompletionToolParamFunctionChunk, tool: object) -> ChatCompletionToolParam: eager_input_streaming: Final = eager_input_streaming_flag(tool) if eager_input_streaming is None: @@ -467,7 +481,11 @@ class LiteLLMAnthropicMessagesAdapter: preserve_midturn_system: bool = False, ) -> list[AllMessageValues]: new_messages: Final[list[AllMessageValues]] = [] - replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages) + replayable_messages: Final = ( + messages + if _keeps_encrypted_reasoning_for_responses(model, custom_llm_provider) + else strip_encrypted_reasoning_blocks_from_anthropic_messages(messages) + ) leading_count: Final = next( (i for i, m in enumerate(replayable_messages) if not is_system_role_message(m)), len(replayable_messages), @@ -1358,9 +1376,9 @@ class LiteLLMAnthropicMessagesAdapter: for choice in choices ): new_content.extend(dict(_COMPACTION_BLOCK.validate_python(block)) for block in compaction_blocks) - # Handle thinking blocks first - if hasattr(choice.message, "thinking_blocks") and choice.message.thinking_blocks: - for thinking_block in choice.message.thinking_blocks: + signed_thinking_blocks = _anthropic_thinking_blocks(choice.message) + if signed_thinking_blocks: + for thinking_block in signed_thinking_blocks: if thinking_block.get("type") == "thinking": if is_empty_unsigned_thinking_block(thinking_block): continue @@ -1627,7 +1645,7 @@ class LiteLLMAnthropicMessagesAdapter: def translate_streaming_openai_chunk_to_anthropic_content_block( self, choices: list[OpenAIStreamingChoice | StreamingChoices] ) -> tuple[ - Literal["text", "tool_use", "thinking"], + Literal["text", "tool_use", "thinking", "redacted_thinking"], "ContentBlockContentBlockDict", ]: from litellm._uuid import uuid @@ -1660,20 +1678,24 @@ class LiteLLMAnthropicMessagesAdapter: choice.delta ) is not None: return "text", TextBlock(type="text", text="") - elif isinstance(choice, StreamingChoices) and hasattr(choice.delta, "thinking_blocks"): - thinking_blocks = choice.delta.thinking_blocks or [] - if len(thinking_blocks) > 0: - thinking_block = thinking_blocks[0] - if thinking_block["type"] == "thinking": - thinking = thinking_block.get("thinking") or "" - signature = thinking_block.get("signature") or "" + elif isinstance(choice, StreamingChoices) and (thinking_blocks := _anthropic_thinking_blocks(choice.delta)): + thinking_block = thinking_blocks[0] + if thinking_block["type"] == "redacted_thinking": + redacted_data = thinking_block.get("data") + return "redacted_thinking", ChatCompletionRedactedThinkingBlock( + type="redacted_thinking", + data=redacted_data if isinstance(redacted_data, str) else "", + ) + if thinking_block["type"] == "thinking": + thinking = thinking_block.get("thinking") or "" + signature = thinking_block.get("signature") or "" - assert isinstance(thinking, str) - assert isinstance(signature, str) + assert isinstance(thinking, str) + assert isinstance(signature, str) - return "thinking", ChatCompletionThinkingBlock( - type="thinking", thinking=thinking, signature=signature - ) + return "thinking", ChatCompletionThinkingBlock( + type="thinking", thinking=thinking, signature=signature + ) # OpenAI-compatible reasoning backends (e.g. vLLM/SGLang reasoning # parsers) populate ``reasoning_content`` without ``thinking_blocks``. # ``Delta`` deletes the ``thinking_blocks`` attribute when unset, so the @@ -1706,19 +1728,17 @@ class LiteLLMAnthropicMessagesAdapter: for tool in choice.delta.tool_calls: if tool.function is not None and tool.function.arguments is not None: partial_json = (partial_json or "") + tool.function.arguments - elif isinstance(choice, StreamingChoices) and hasattr(choice.delta, "thinking_blocks"): - thinking_blocks = choice.delta.thinking_blocks or [] - if len(thinking_blocks) > 0: - for thinking_block in thinking_blocks: - if thinking_block["type"] == "thinking": - thinking = thinking_block.get("thinking") or "" - signature = thinking_block.get("signature") or "" + elif isinstance(choice, StreamingChoices) and (thinking_blocks := _anthropic_thinking_blocks(choice.delta)): + for thinking_block in thinking_blocks: + if thinking_block["type"] == "thinking": + thinking = thinking_block.get("thinking") or "" + signature = thinking_block.get("signature") or "" - assert isinstance(thinking, str) - assert isinstance(signature, str) + assert isinstance(thinking, str) + assert isinstance(signature, str) - reasoning_content += thinking - reasoning_signature += signature + reasoning_content += thinking + reasoning_signature += signature # Handle reasoning_content when thinking_blocks is not present # This handles providers like OpenRouter that return reasoning_content elif isinstance(choice, StreamingChoices) and hasattr(choice.delta, "reasoning_content"): diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index d5e3dfc0c02..7ed3ceda1eb 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -4,7 +4,10 @@ from urllib.parse import urlsplit import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_llm_provider, + resolve_model_and_provider_for_metadata, +) from litellm.llms.azure_ai.chat.transformation import AzureAIGPT5Config from litellm.llms.azure_ai.common_utils import ( azure_ai_supports_native_responses, @@ -186,3 +189,23 @@ def responses_api_bridge_check( model = model.replace("responses/", "") return model_info, model + + +def chat_completion_ends_in_responses_api(model: str | None, custom_llm_provider: str | None) -> bool: + """Whether chat completion for this model is served by a Responses API. + + Encrypted reasoning can be replayed only there. A normal chat provider cannot + verify the block, so the messages adapter strips it. ``chatgpt`` and a + ``responses/`` model keep it. + """ + if model is None and custom_llm_provider is None: + return False + try: + resolved_model, provider = resolve_model_and_provider_for_metadata(model or "", custom_llm_provider) + except Exception as error: # noqa: BLE001 # unknown model names must not fail message translation + verbose_logger.debug("responses replay check failed for %s: %s", model, error) + return custom_llm_provider == LlmProviders.CHATGPT.value or ( + isinstance(model, str) and (model == "chatgpt" or model.startswith("chatgpt/")) + ) + model_info, _resolved = responses_api_bridge_check(model=resolved_model, custom_llm_provider=provider) + return model_info.get("mode") == "responses" diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index eb308a2c513..f24f163541f 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3240,6 +3240,165 @@ def test_reasoning_items_streaming_emitted_on_response_completed(): assert ri["summary"][0]["text"] == summary_text +def test_reasoning_output_item_done_emits_signed_thinking_blocks_the_stream_keeps(): + """A chatgpt /v1/messages stream drops the reasoning item unless this delta carries thinking_blocks. + + response.output_item.done is the only event that has encrypted_content. The chat delta has to + expose it as a signed thinking block, and that block alone has to count as stream content. + """ + from unittest.mock import MagicMock + + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter + + encrypted: Final = "gAAAA_round_trip" + done: Final = { + "type": "response.output_item.done", + "output_index": 0, + "item": { + "type": "reasoning", + "id": "rs_done", + "encrypted_content": encrypted, + "summary": [], + }, + } + result: Final = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True + ).chunk_parser(done) + blocks: Final = result.choices[0].delta.thinking_blocks + assert blocks is not None + assert blocks[0]["type"] == "redacted_thinking" + assert blocks[0]["data"] == encrypted_reasoning_signature(encrypted) + assert blocks[0]["index"] == 0 + assert getattr(result.choices[0].delta, "reasoning_content", None) in (None, "") + + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model="chatgpt/gpt-6.1-sol", + logging_obj=MagicMock(), + custom_llm_provider="chatgpt", + ) + assert wrapper.is_chunk_non_empty({}, result, {}) is True + + block_type, start = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_chunk_to_anthropic_content_block( + choices=result.choices + ) + assert block_type == "redacted_thinking" + assert start["data"] == encrypted_reasoning_signature(encrypted) + + +def test_reasoning_output_item_done_with_summary_signs_the_thinking_text(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature + + encrypted: Final = "gAAAA_with_summary" + done: Final = { + "type": "response.output_item.done", + "output_index": 2, + "item": { + "type": "reasoning", + "id": "rs_sum", + "encrypted_content": encrypted, + "summary": [{"type": "summary_text", "text": "Oslo is colder."}], + }, + } + result: Final = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True + ).chunk_parser(done) + block: Final = result.choices[0].delta.thinking_blocks[0] + assert block["type"] == "thinking" + assert block["thinking"] == "Oslo is colder." + assert block["signature"] == encrypted_reasoning_signature(encrypted) + assert block["index"] == 2 + + +def test_reasoning_output_item_done_without_encrypted_content_stays_empty(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + done: Final = { + "type": "response.output_item.done", + "output_index": 0, + "item": { + "type": "reasoning", + "id": "rs_plain", + "summary": [{"type": "summary_text", "text": "already streamed"}], + }, + } + result: Final = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True + ).chunk_parser(done) + assert getattr(result.choices[0].delta, "thinking_blocks", None) in (None, []) + + +def test_transform_response_copies_encrypted_reasoning_onto_thinking_blocks(): + from unittest.mock import Mock + + from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.utils import ModelResponse, Usage + + encrypted: Final = "gAAAA_non_stream" + reasoning: Final = ResponseReasoningItem( + id="rs_tool", + summary=[], + type="reasoning", + content=None, + encrypted_content=encrypted, + status=None, + ) + function_call: Final = ResponseFunctionToolCall( + id="fc_1", + type="function_call", + status="completed", + arguments='{"city":"Oslo"}', + call_id="call_1", + name="get_weather", + ) + raw: Final = ResponsesAPIResponse.model_construct( + id="resp_tool", + created_at=1, + error=None, + incomplete_details=None, + model="gpt-6.1-sol", + output=[reasoning, function_call], + status="completed", + usage=None, + ) + model_response: Final = ModelResponse( + id="chatcmpl-tool", + created=1, + model=None, + object="chat.completion", + system_fingerprint=None, + choices=[], + usage=Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0), + ) + result: Final = LiteLLMResponsesTransformationHandler().transform_response( + model="gpt-6.1-sol", + raw_response=raw, + model_response=model_response, + logging_obj=Mock(), + request_data={"model": "gpt-6.1-sol"}, + messages=[{"role": "user", "content": "Which city?"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + message: Final = result.choices[0].message + assert message.reasoning_items[0]["encrypted_content"] == encrypted + assert message.thinking_blocks[0]["type"] == "redacted_thinking" + assert message.thinking_blocks[0]["data"] == encrypted_reasoning_signature(encrypted) + assert message.tool_calls[0].function.name == "get_weather" + + def test_streaming_function_call_tool_id_for_degenerate_call_id(): """In streaming, Bedrock Mantle's function_call event carries a unique ``id`` (``fc_...``) and a non-unique, index-based ``call_id`` (``call_0``). For that diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 637811f6387..8dfce043fe1 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -246,6 +246,37 @@ def test_is_chunk_non_empty(initialized_custom_stream_wrapper: CustomStreamWrapp ) +def test_is_chunk_non_empty_keeps_thinking_blocks_without_other_content( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """Encrypted reasoning arrives as thinking_blocks and nothing else, so that delta has to be yielded.""" + chunk: Final = ModelResponseStream( + id="chunk-thinking", + object="chat.completion.chunk", + created=1, + model="chatgpt/gpt-6.1-sol", + choices=[ + StreamingChoices( + index=0, + finish_reason=None, + delta=Delta( + thinking_blocks=[ + {"type": "redacted_thinking", "data": "litellm_encrypted_reasoning:gAAAA"} + ] + ), + ) + ], + ) + assert ( + initialized_custom_stream_wrapper.is_chunk_non_empty( + completion_obj={}, + model_response=chunk, + response_obj={}, + ) + is True + ) + + def test_is_chunk_non_empty_with_annotations( initialized_custom_stream_wrapper: CustomStreamWrapper, ): diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 2f0a7ad070b..49a8341ce9e 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -509,6 +509,166 @@ def test_translate_anthropic_messages_to_openai_drops_bridge_encrypted_reasoning assert [b["signature"] for b in result[2]["thinking_blocks"]] == ["EqQBCkYIAxgCIkA_signed"] +def test_chatgpt_messages_keep_encrypted_reasoning_and_replay_it(): + """chatgpt /v1/messages ends in the Responses API, so the signed block has to survive into the next request.""" + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + summary_encrypted: Final = "gAAAA_summary" + bytes_encrypted: Final = "gAAAA_bytes" + messages: Final = [ + AnthropicMessagesUserMessageParam(role="user", content="Which city is colder?"), + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[ + { + "type": "thinking", + "thinking": "Oslo is colder.", + "signature": encrypted_reasoning_signature(summary_encrypted), + }, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature(bytes_encrypted)}, + {"type": "tool_use", "id": "call_1", "name": "get_weather", "input": {"city": "Oslo"}}, + ], + ), + AnthropicMessagesUserMessageParam( + role="user", + content=[{"type": "tool_result", "tool_use_id": "call_1", "content": "-7 C"}], + ), + ] + + kept: Final = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=messages, + model="chatgpt/gpt-6.1-sol", + custom_llm_provider="chatgpt", + ) + assistant: Final = kept[1] + assert [block["type"] for block in assistant["thinking_blocks"]] == ["thinking", "redacted_thinking"] + assert assistant["thinking_blocks"][0]["signature"] == encrypted_reasoning_signature(summary_encrypted) + assert assistant["thinking_blocks"][1]["data"] == encrypted_reasoning_signature(bytes_encrypted) + + responses_input, _instructions = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api( + kept + ) + reasoning: Final = tuple(item for item in responses_input if item.get("type") == "reasoning") + assert [item["encrypted_content"] for item in reasoning] == [summary_encrypted, bytes_encrypted] + types: Final = [item.get("type") for item in responses_input] + assert types.index("reasoning") < types.index("function_call") + + stripped: Final = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=messages, + model="gpt-4o", + custom_llm_provider="openai", + ) + assert not stripped[1].get("thinking_blocks") + + prefixed: Final = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=messages, + model="responses/gpt-5-mini", + custom_llm_provider="openai", + ) + assert prefixed[1]["thinking_blocks"][1]["data"] == encrypted_reasoning_signature(bytes_encrypted) + + +def test_translate_openai_reasoning_items_become_signed_thinking_blocks(): + encrypted: Final = "gAAAA_outbound" + openai_choices: Final = [ + Choices( + message=Message( + role="assistant", + content="Oslo.", + reasoning_items=[ + { + "type": "reasoning", + "id": "rs_1", + "encrypted_content": encrypted, + "summary": [{"type": "summary_text", "text": "Oslo is colder."}], + } + ], + tool_calls=[ + ChatCompletionAssistantToolCall( + id="call_1", + type="function", + function=Function(name="get_weather", arguments='{"city": "Oslo"}'), + ) + ], + ) + ) + ] + + result: Final = LiteLLMAnthropicMessagesAdapter()._translate_openai_content_to_anthropic(choices=openai_choices) + + assert result[0]["type"] == "thinking" + assert result[0]["thinking"] == "Oslo is colder." + assert result[0]["signature"] == encrypted_reasoning_signature(encrypted) + assert result[1]["type"] == "text" + assert result[2]["type"] == "tool_use" + + +def test_translate_openai_reasoning_items_without_summary_become_redacted_thinking(): + encrypted: Final = "gAAAA_bytes_only" + openai_choices: Final = [ + Choices( + message=Message( + role="assistant", + content=None, + reasoning_items=[ + {"type": "reasoning", "id": "rs_1", "encrypted_content": encrypted, "summary": []} + ], + ) + ) + ] + + result: Final = LiteLLMAnthropicMessagesAdapter()._translate_openai_content_to_anthropic(choices=openai_choices) + + assert result == [ + {"type": "redacted_thinking", "data": encrypted_reasoning_signature(encrypted)}, + ] + + +def test_translate_openai_reasoning_items_are_not_duplicated_when_thinking_blocks_already_carry_them(): + encrypted: Final = "gAAAA_once" + signature: Final = encrypted_reasoning_signature(encrypted) + openai_choices: Final = [ + Choices( + message=Message( + role="assistant", + content="Oslo.", + thinking_blocks=[{"type": "redacted_thinking", "data": signature}], + reasoning_items=[ + {"type": "reasoning", "id": "rs_1", "encrypted_content": encrypted, "summary": []} + ], + ) + ) + ] + + result: Final = LiteLLMAnthropicMessagesAdapter()._translate_openai_content_to_anthropic(choices=openai_choices) + + assert [block["type"] for block in result] == ["redacted_thinking", "text"] + + +def test_streaming_chunk_with_only_reasoning_items_opens_a_signed_thinking_block(): + encrypted: Final = "gAAAA_delta" + choices: Final = [ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_items=[ + {"type": "reasoning", "id": "rs_1", "encrypted_content": encrypted, "summary": []} + ] + ), + ) + ] + + block_type, start = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_chunk_to_anthropic_content_block( + choices=choices + ) + + assert block_type == "redacted_thinking" + assert start["data"] == encrypted_reasoning_signature(encrypted) + + def test_translate_anthropic_messages_to_openai_sets_reasoning_content(): """Reasoning-aware chat providers read reasoning_content, so thinking text must land there. From 6fef70ac726abf7efee91e6b72ad262a81fb7380 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 13:17:32 -0800 Subject: [PATCH 05/41] fix(anthropic): keep Responses reasoning on tool turns and retry once when a backend cannot decrypt it Backport of the second #49076 commit to lts/1.108.x. Cherry-picked from 10c953d8556. Adapted to the line: chat_completion_ends_in_responses_api has no deployment_model_info parameter here, so only tools, reasoning_effort, and api_base are threaded through, and the moved tests import the adapters from pass_through instead of messages. --- .../transformation.py | 35 ++- .../prompt_templates/common_utils.py | 8 + .../pass_through/adapters/handler.py | 97 +++++++-- .../adapters/streaming_iterator.py | 8 - .../pass_through/adapters/transformation.py | 69 ++++-- litellm/responses/bridge_check.py | 21 +- ...responses_transformation_transformation.py | 85 +++++++- ...al_pass_through_adapters_transformation.py | 205 ++++++++++++++++-- 8 files changed, 443 insertions(+), 85 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 7fe36c2e42a..296f1949d11 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -448,6 +448,7 @@ def _as_chat_reasoning_items( def _with_stream_block_index( block: Mapping[str, object], index: int, + summary_streamed: bool, ) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock: if block.get("type") == "redacted_thinking": data: Final = block.get("data") @@ -460,21 +461,33 @@ def _with_stream_block_index( signature: Final = block.get("signature") return ChatCompletionThinkingBlock( type="thinking", - thinking=thinking if isinstance(thinking, str) else "", + thinking=thinking if isinstance(thinking, str) and not summary_streamed else "", signature=signature if isinstance(signature, str) else None, index=index, ) +def _stream_output_index(chunk: Mapping[str, object]) -> int | None: + output_index: Final = chunk.get("output_index") + return output_index if isinstance(output_index, int) and not isinstance(output_index, bool) else None + + def _reasoning_done_thinking_blocks( item: Mapping[str, object], - output_index: object, + output_index: int | None, + summary_streamed: bool, ) -> tuple[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, ...]: - index: Final = output_index if isinstance(output_index, int) and not isinstance(output_index, bool) else 0 - return tuple(_with_stream_block_index(block, index) for block in thinking_blocks_from_reasoning_items((item,))) + index: Final = output_index if output_index is not None else 0 + return tuple( + _with_stream_block_index(block, index, summary_streamed) + for block in thinking_blocks_from_reasoning_items((item,)) + ) -def _reasoning_done_stream_chunk(chunk: Mapping[str, object]) -> ModelResponseStream | None: +def _reasoning_done_stream_chunk( + chunk: Mapping[str, object], + summary_streamed_outputs: frozenset[int | None], +) -> ModelResponseStream | None: event_type: Final = chunk.get("type") event_name: Final = ( event_type.value @@ -488,7 +501,8 @@ def _reasoning_done_stream_chunk(chunk: Mapping[str, object]) -> ModelResponseSt item: Final = as_str_mapping(chunk.get("item")) if item is None or item.get("type") != "reasoning": return None - signed_blocks: Final = _reasoning_done_thinking_blocks(item, chunk.get("output_index")) + output_index: Final = _stream_output_index(chunk) + signed_blocks: Final = _reasoning_done_thinking_blocks(item, output_index, output_index in summary_streamed_outputs) if not signed_blocks: return None return ModelResponseStream( @@ -556,7 +570,6 @@ def _fold_choices_into_one( *chain.from_iterable(_message_reasoning_items(m) for m in messages), *(() if trailing_reasoning_item is None else (trailing_reasoning_item,)), ) - encrypted_thinking_blocks: Final = thinking_blocks_from_reasoning_items(reasoning_items) return Choices( message=Message( role=messages[0].role if messages else "assistant", @@ -565,7 +578,6 @@ def _fold_choices_into_one( tool_calls=list(tool_calls) or None, reasoning_content=" ".join(value for value in reasoning_contents if value) or None, reasoning_items=_as_chat_reasoning_items(reasoning_items), - thinking_blocks=list(encrypted_thinking_blocks) or None, provider_specific_fields=_tool_search_provider_fields(tool_search_items), ), finish_reason="tool_calls" if tool_calls else "stop", @@ -1371,12 +1383,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): for summary_block in reasoning_item["summary"] if summary_block.get("text") ) - encrypted_thinking_blocks: Final = thinking_blocks_from_reasoning_items(reasoning_items) message: Final = Message( content=None if tool_search_items else "", reasoning_content=reasoning_content if reasoning_content else None, reasoning_items=_as_chat_reasoning_items(reasoning_items), - thinking_blocks=list(encrypted_thinking_blocks) or None, provider_specific_fields=_tool_search_provider_fields(tool_search_items), ) return Choices(message=message, finish_reason=finish_reason, index=0) @@ -1941,6 +1951,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): self._tool_call_index_map: dict[int, int] = {} # mutable-ok: per-stream accumulator state self._structured_output: Final = structured_output self._first_message_item_done: bool = False + self._summary_streamed_outputs: frozenset[int | None] = frozenset() def _handle_string_chunk( self, str_line: Union[str, "BaseModel"] @@ -2255,7 +2266,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): """ verbose_logger.debug("Chat provider: transform_streaming_response called with chunk: %s", chunk) self._remember_served_service_tier(chunk) - reasoning_chunk: Final = _reasoning_done_stream_chunk(chunk) + reasoning_chunk: Final = _reasoning_done_stream_chunk(chunk, self._summary_streamed_outputs) translated: Final = ( self._empty_content_chunk() if self._is_text_delta_after_structured_answer(chunk) @@ -2267,6 +2278,8 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ) if self._structured_output and _is_message_item_done(chunk): self._first_message_item_done = True + if _responses_event_type(chunk) == "response.reasoning_summary_text.delta" and chunk.get("delta"): + self._summary_streamed_outputs = self._summary_streamed_outputs | {_stream_output_index(chunk)} return self._with_served_service_tier(self._with_stream_scoped_id(translated)) def _is_text_delta_after_structured_answer(self, chunk: Mapping[str, object]) -> bool: diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 790860ca561..f5fa6f649a6 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -2462,6 +2462,14 @@ def anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]: ) +def _content_carries_encrypted_reasoning(content: object) -> bool: + return any(is_encrypted_reasoning_block(block) for block in _object_sequence(content)) + + +def messages_carry_encrypted_reasoning(messages: Sequence[object]) -> bool: + return any(_content_carries_encrypted_reasoning(content) for content in anthropic_content_lists(messages)) + + def _strip_encrypted_reasoning_from_blocks( content: object, *, diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 8c0c065c457..fda9fe61060 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -1,5 +1,6 @@ -from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Mapping, Sequence from dataclasses import dataclass +from functools import partial from typing import ( TYPE_CHECKING, Final, @@ -13,12 +14,15 @@ from typing_extensions import TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.litellm_core_utils.prompt_templates.common_utils import messages_carry_encrypted_reasoning from litellm.litellm_core_utils.tool_fields import ( ROUTES_WITHOUT_MCP_SERVERS, drops_mcp_servers, tools_without_mcp_servers, ) +from litellm.llms.anthropic.common_utils import strip_encrypted_reasoning_blocks_from_anthropic_messages from litellm.llms.anthropic.pass_through.adapters.transformation import ( AnthropicAdapter, ) @@ -51,6 +55,7 @@ from litellm.types.utils import ModelResponse from litellm.utils import get_model_info if TYPE_CHECKING: + from litellm.chat_completions.dispatch import ChatResult from litellm.router import Router from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth @@ -68,22 +73,68 @@ class _CompletionKwargs(TypedDict, total=False, extra_items=object): custom_llm_provider: str +_ANTHROPIC_MESSAGES: Final = TypeAdapter(list[dict[str, object]]) +_PrepareReplay: TypeAlias = "Callable[[_AnthropicMessages], _CompletionKwargs]" + + @dataclass(frozen=True, slots=True) class _ChatTargetReasoning: effort: str summary: str | None = None +def _messages_without_rejected_reasoning( + messages: _AnthropicMessages, + error: litellm.BadRequestError, +) -> _AnthropicMessages | None: + error_message: Final = str(error) + rejected: Final = "invalid_encrypted_content" in error_message or "could not be verified" in error_message + if not rejected or not messages_carry_encrypted_reasoning(messages): + return None + verbose_logger.warning("Backend rejected replayed encrypted reasoning, retrying once without it: %s", error) + return _ANTHROPIC_MESSAGES.validate_python(strip_encrypted_reasoning_blocks_from_anthropic_messages(messages)) + + +async def _acompletion_without_rejected_reasoning( + completion_kwargs: _CompletionKwargs, + messages: _AnthropicMessages, + prepare_replay: _PrepareReplay, +) -> tuple[_CompletionKwargs, "ChatResult"]: + try: + return completion_kwargs, await litellm.acompletion(**completion_kwargs) + except litellm.BadRequestError as error: + replayed_messages: Final = _messages_without_rejected_reasoning(messages, error) + if replayed_messages is None: + raise + replay_kwargs: Final = prepare_replay(replayed_messages) + return replay_kwargs, await litellm.acompletion(**replay_kwargs) + + +def _completion_without_rejected_reasoning( + completion_kwargs: _CompletionKwargs, + messages: _AnthropicMessages, + prepare_replay: _PrepareReplay, +) -> tuple[_CompletionKwargs, "ChatResult | Coroutine[object, object, ChatResult]"]: + try: + return completion_kwargs, litellm.completion(**completion_kwargs) + except litellm.BadRequestError as error: + replayed_messages: Final = _messages_without_rejected_reasoning(messages, error) + if replayed_messages is None: + raise + replay_kwargs: Final = prepare_replay(replayed_messages) + return replay_kwargs, litellm.completion(**replay_kwargs) + + +def _is_compaction_block(block: object) -> bool: + mapping: Final = as_str_mapping(block) + return mapping is not None and mapping.get("type") == "compaction" + + def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool: """Return True when any message carries a ``compaction`` content block.""" - for msg in messages: - content = msg.get("content") - if not isinstance(content, list): - continue - for block in content: - if isinstance(block, dict) and block.get("type") == "compaction": - return True - return False + return any( + any(map(_is_compaction_block, content)) for msg in messages if isinstance(content := msg.get("content"), list) + ) def _proxy_router_fallback() -> "Router | None": @@ -646,12 +697,14 @@ class LiteLLMMessagesToCompletionTransformationHandler: if "output_config" in extra_kwargs: request_data["output_config"] = extra_kwargs["output_config"] + api_base: Final = provider_kwargs.get("api_base") ( openai_request, tool_name_mapping, ) = ANTHROPIC_ADAPTER.translate_completion_input_params_with_tool_mapping( request_data, custom_llm_provider=target_provider, + api_base=api_base if isinstance(api_base, str) else None, ) if openai_request is None: @@ -779,12 +832,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: parent_trace_id=sub_call_trace_id, ) - ( - completion_kwargs, - tool_name_mapping, - ) = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( + prepare: Final = partial( + LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs, max_tokens=max_tokens, - messages=effective_messages, model=model, metadata=metadata, stop_sequences=stop_sequences, @@ -799,8 +849,10 @@ class LiteLLMMessagesToCompletionTransformationHandler: output_format=output_format, extra_kwargs=kwargs, ) - - completion_response: Final = await litellm.acompletion(**completion_kwargs) + initial_kwargs, tool_name_mapping = prepare(messages=effective_messages) + completion_kwargs, completion_response = await _acompletion_without_rejected_reasoning( + initial_kwargs, effective_messages, lambda replayed_messages: prepare(messages=replayed_messages)[0] + ) logging_obj: Final = litellm_logging_obj_from_kwargs(kwargs) if stream: @@ -924,12 +976,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: effective_messages: Final = polyfill_result.messages if polyfill_result is not None else messages effective_system: Final = polyfill_result.system if polyfill_result is not None else system - ( - completion_kwargs, - tool_name_mapping, - ) = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( + prepare: Final = partial( + LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs, max_tokens=max_tokens, - messages=effective_messages, model=model, metadata=metadata, stop_sequences=stop_sequences, @@ -944,8 +993,10 @@ class LiteLLMMessagesToCompletionTransformationHandler: output_format=output_format, extra_kwargs=kwargs, ) - - completion_response: Final = litellm.completion(**completion_kwargs) + initial_kwargs, tool_name_mapping = prepare(messages=effective_messages) + completion_kwargs, completion_response = _completion_without_rejected_reasoning( + initial_kwargs, effective_messages, lambda replayed_messages: prepare(messages=replayed_messages)[0] + ) logging_obj: Final = litellm_logging_obj_from_kwargs(kwargs) if stream: diff --git a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py index 0f1398ac9b2..cc705ab1184 100644 --- a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py @@ -24,9 +24,7 @@ from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.exceptions import ContextWindowExceededError, MidStreamFallbackError -from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.litellm_core_utils.hidden_params import set_hidden_params -from litellm.litellm_core_utils.prompt_templates.common_utils import thinking_blocks_from_reasoning_items from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -68,10 +66,6 @@ def _optional_attr(obj: object, name: str) -> object: return getattr(obj, name, None) -def _delta_has_encrypted_reasoning(delta: object) -> bool: - return bool(thinking_blocks_from_reasoning_items(attribute_or_none(delta, "reasoning_items"))) - - def _optional_attr_sequence(obj: object, name: str) -> Sequence[object]: value: Final = getattr(obj, name, None) return value if value else () @@ -1312,8 +1306,6 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): isinstance(b, dict) and not is_empty_unsigned_thinking_block(b) for b in thinking_blocks ): return False - if _delta_has_encrypted_reasoning(delta): - return False return True def _should_start_new_content_block(self, chunk: "ModelResponseStream") -> bool: diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index 62f33bb11fd..ae7c433de2a 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -2,12 +2,14 @@ import copy import hashlib import json from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast from pydantic import JsonValue, TypeAdapter import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_items from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.llms.anthropic.pass_through.utils import ( is_reasoning_auto_summary_enabled, @@ -220,10 +222,30 @@ def target_supports_mid_conversation_system(model: str | None, custom_llm_provid return supports_mid_conversation_system(model=model, custom_llm_provider=custom_llm_provider) -def _keeps_encrypted_reasoning_for_responses(model: str | None, custom_llm_provider: str | None) -> bool: +@dataclass(frozen=True, slots=True) +class ChatRouteInputs: + tools: Sequence[Mapping[str, object]] | None = None + reasoning_effort: str | Mapping[str, object] | None = None + api_base: str | None = None + + +_NO_ROUTE_INPUTS: Final = ChatRouteInputs() + + +def _keeps_encrypted_reasoning_for_responses( + model: str | None, + custom_llm_provider: str | None, + route: ChatRouteInputs, +) -> bool: from litellm.responses.bridge_check import chat_completion_ends_in_responses_api - return chat_completion_ends_in_responses_api(model, custom_llm_provider) + return chat_completion_ends_in_responses_api( + model, + custom_llm_provider, + tools=route.tools, + reasoning_effort=route.reasoning_effort, + api_base=route.api_base, + ) def _anthropic_thinking_blocks(source: object) -> tuple[Mapping[str, object], ...]: @@ -233,6 +255,10 @@ def _anthropic_thinking_blocks(source: object) -> tuple[Mapping[str, object], .. ) +def _streamed_thinking_blocks(delta: object) -> tuple[Mapping[str, object], ...]: + return tuple(block for item in attribute_items(delta, "thinking_blocks") if (block := as_str_mapping(item))) + + def _chat_tool_param(function_chunk: ChatCompletionToolParamFunctionChunk, tool: object) -> ChatCompletionToolParam: eager_input_streaming: Final = eager_input_streaming_flag(tool) if eager_input_streaming is None: @@ -260,7 +286,7 @@ class AnthropicAdapter: return result def translate_completion_input_params_with_tool_mapping( - self, kwargs, *, custom_llm_provider: str | None = None + self, kwargs, *, custom_llm_provider: str | None = None, api_base: str | None = None ) -> tuple[ChatCompletionRequest | None, dict[str, str]]: """ Translate Anthropic request params to OpenAI format, returning tool name mapping. @@ -294,6 +320,7 @@ class AnthropicAdapter: ) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( anthropic_message_request=request_body, custom_llm_provider=custom_llm_provider, + api_base=api_base, ) return translated_body, tool_name_mapping @@ -479,11 +506,12 @@ class LiteLLMAnthropicMessagesAdapter: *, custom_llm_provider: str | None = None, preserve_midturn_system: bool = False, + route: ChatRouteInputs = _NO_ROUTE_INPUTS, ) -> list[AllMessageValues]: new_messages: Final[list[AllMessageValues]] = [] replayable_messages: Final = ( messages - if _keeps_encrypted_reasoning_for_responses(model, custom_llm_provider) + if _keeps_encrypted_reasoning_for_responses(model, custom_llm_provider, route) else strip_encrypted_reasoning_blocks_from_anthropic_messages(messages) ) leading_count: Final = next( @@ -1239,6 +1267,7 @@ class LiteLLMAnthropicMessagesAdapter: anthropic_message_request: AnthropicMessagesRequest, *, custom_llm_provider: str | None = None, + api_base: str | None = None, preserve_midturn_system: bool = False, ) -> tuple[ChatCompletionRequest, dict[str, str]]: """ @@ -1250,26 +1279,15 @@ class LiteLLMAnthropicMessagesAdapter: for tools that exceeded OpenAI's 64-char limit """ # Debug: Processing Anthropic message request - new_messages: list[AllMessageValues] = [] tool_name_mapping: dict[str, str] = {} - ## CONVERT ANTHROPIC MESSAGES TO OPENAI messages_list: Final[list[AllAnthropicPassThroughMessageValues]] = cast( list[AllAnthropicPassThroughMessageValues], anthropic_message_request["messages"], ) - new_messages = self.translate_anthropic_messages_to_openai( - messages=messages_list, - model=anthropic_message_request.get("model"), - custom_llm_provider=custom_llm_provider, - preserve_midturn_system=preserve_midturn_system, - ) - ## ADD SYSTEM MESSAGE TO MESSAGES - self._add_system_message_to_messages(new_messages, anthropic_message_request) - new_kwargs: Final[ChatCompletionRequest] = { "model": anthropic_message_request["model"], - "messages": new_messages, + "messages": [], } ## CONVERT METADATA (user_id + litellm metadata) self._translate_metadata_to_openai( @@ -1293,6 +1311,21 @@ class LiteLLMAnthropicMessagesAdapter: new_kwargs=new_kwargs, custom_llm_provider=custom_llm_provider, ) + ## CONVERT ANTHROPIC MESSAGES TO OPENAI, after tools and thinking pick the route + new_messages: Final = self.translate_anthropic_messages_to_openai( + messages=messages_list, + model=anthropic_message_request.get("model"), + custom_llm_provider=custom_llm_provider, + preserve_midturn_system=preserve_midturn_system, + route=ChatRouteInputs( + tools=new_kwargs.get("tools"), + reasoning_effort=new_kwargs.get("reasoning_effort"), + api_base=api_base, + ), + ) + ## ADD SYSTEM MESSAGE TO MESSAGES + self._add_system_message_to_messages(new_messages, anthropic_message_request) + new_kwargs["messages"] = new_messages ## CONVERT STOP_SEQUENCES self._translate_stop_sequences_to_openai( anthropic_message_request=anthropic_message_request, @@ -1678,7 +1711,7 @@ class LiteLLMAnthropicMessagesAdapter: choice.delta ) is not None: return "text", TextBlock(type="text", text="") - elif isinstance(choice, StreamingChoices) and (thinking_blocks := _anthropic_thinking_blocks(choice.delta)): + elif isinstance(choice, StreamingChoices) and (thinking_blocks := _streamed_thinking_blocks(choice.delta)): thinking_block = thinking_blocks[0] if thinking_block["type"] == "redacted_thinking": redacted_data = thinking_block.get("data") @@ -1728,7 +1761,7 @@ class LiteLLMAnthropicMessagesAdapter: for tool in choice.delta.tool_calls: if tool.function is not None and tool.function.arguments is not None: partial_json = (partial_json or "") + tool.function.arguments - elif isinstance(choice, StreamingChoices) and (thinking_blocks := _anthropic_thinking_blocks(choice.delta)): + elif isinstance(choice, StreamingChoices) and (thinking_blocks := _streamed_thinking_blocks(choice.delta)): for thinking_block in thinking_blocks: if thinking_block["type"] == "thinking": thinking = thinking_block.get("thinking") or "" diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index 7ed3ceda1eb..cdd10d489a1 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -191,12 +191,19 @@ def responses_api_bridge_check( return model_info, model -def chat_completion_ends_in_responses_api(model: str | None, custom_llm_provider: str | None) -> bool: +def chat_completion_ends_in_responses_api( + model: str | None, + custom_llm_provider: str | None, + *, + tools: Sequence[Mapping[str, object]] | None = None, + reasoning_effort: str | Mapping[str, object] | None = None, + api_base: str | None = None, +) -> bool: """Whether chat completion for this model is served by a Responses API. Encrypted reasoning can be replayed only there. A normal chat provider cannot - verify the block, so the messages adapter strips it. ``chatgpt`` and a - ``responses/`` model keep it. + verify the block, so the messages adapter strips it. ``chatgpt``, a ``responses/`` + model, and a request whose tools or reasoning effort the bridge sends to Responses keep it. """ if model is None and custom_llm_provider is None: return False @@ -207,5 +214,11 @@ def chat_completion_ends_in_responses_api(model: str | None, custom_llm_provider return custom_llm_provider == LlmProviders.CHATGPT.value or ( isinstance(model, str) and (model == "chatgpt" or model.startswith("chatgpt/")) ) - model_info, _resolved = responses_api_bridge_check(model=resolved_model, custom_llm_provider=provider) + model_info, _resolved = responses_api_bridge_check( + model=resolved_model, + custom_llm_provider=provider, + tools=tools, + reasoning_effort=reasoning_effort, + api_base=api_base, + ) return model_info.get("mode") == "responses" diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index f24f163541f..d1155e41bdf 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3253,7 +3253,7 @@ def test_reasoning_output_item_done_emits_signed_thinking_blocks_the_stream_keep ) from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter + from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter encrypted: Final = "gAAAA_round_trip" done: Final = { @@ -3318,6 +3318,74 @@ def test_reasoning_output_item_done_with_summary_signs_the_thinking_text(): assert block["index"] == 2 +def test_streamed_reasoning_summary_reaches_messages_clients_once_with_its_signature(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature + from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import AnthropicStreamWrapper + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + encrypted: Final = "gAAAA_streamed_summary" + events: Final = ( + {"type": "response.reasoning_summary_part.added", "output_index": 0, "summary_index": 0}, + {"type": "response.reasoning_summary_text.delta", "output_index": 0, "delta": "Oslo is "}, + {"type": "response.reasoning_summary_text.delta", "output_index": 0, "delta": "colder."}, + { + "type": "response.output_item.done", + "output_index": 0, + "item": { + "type": "reasoning", + "id": "rs_streamed", + "encrypted_content": encrypted, + "summary": [{"type": "summary_text", "text": "Oslo is colder."}], + }, + }, + {"type": "response.output_text.delta", "output_index": 1, "delta": "Oslo."}, + ) + parser: Final = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + chunks: Final = ( + *(parser.chunk_parser(event) for event in events), + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")]), + ) + + sse_events: Final = list(AnthropicStreamWrapper(completion_stream=iter(chunks), model="chatgpt/gpt-5.5")) + deltas: Final = [event["delta"] for event in sse_events if event.get("type") == "content_block_delta"] + assert "".join(d.get("thinking", "") for d in deltas if d["type"] == "thinking_delta") == "Oslo is colder." + assert [d["signature"] for d in deltas if d["type"] == "signature_delta"] == [ + encrypted_reasoning_signature(encrypted) + ] + + +def test_reasoning_item_repeated_at_stream_end_reaches_messages_clients_once(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature + from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import AnthropicStreamWrapper + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + encrypted: Final = "gAAAA_completed_repeat" + reasoning_item: Final = {"type": "reasoning", "id": "rs_repeat", "summary": [], "encrypted_content": encrypted} + events: Final = ( + {"type": "response.output_item.done", "output_index": 0, "item": reasoning_item}, + {"type": "response.output_text.delta", "output_index": 1, "delta": "Oslo."}, + ) + parser: Final = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + chunks: Final = ( + *(parser.chunk_parser(event) for event in events), + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(reasoning_items=[reasoning_item]), finish_reason=None)] + ), + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")]), + ) + + sse_events: Final = list(AnthropicStreamWrapper(completion_stream=iter(chunks), model="azure/gpt-5.4-nano")) + started: Final = [event["content_block"] for event in sse_events if event.get("type") == "content_block_start"] + assert [block["type"] for block in started] == ["redacted_thinking", "text"] + assert started[0]["data"] == encrypted_reasoning_signature(encrypted) + + def test_reasoning_output_item_done_without_encrypted_content_stays_empty(): from litellm.completion_extras.litellm_responses_transformation.transformation import ( OpenAiResponsesToChatCompletionStreamIterator, @@ -3338,10 +3406,13 @@ def test_reasoning_output_item_done_without_encrypted_content_stays_empty(): assert getattr(result.choices[0].delta, "thinking_blocks", None) in (None, []) -def test_transform_response_copies_encrypted_reasoning_onto_thinking_blocks(): +def test_transform_response_hands_encrypted_reasoning_to_messages_clients_only(): from unittest.mock import Mock from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature + from litellm.llms.anthropic.pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, + ) from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ModelResponse, Usage @@ -3394,10 +3465,16 @@ def test_transform_response_copies_encrypted_reasoning_onto_thinking_blocks(): ) message: Final = result.choices[0].message assert message.reasoning_items[0]["encrypted_content"] == encrypted - assert message.thinking_blocks[0]["type"] == "redacted_thinking" - assert message.thinking_blocks[0]["data"] == encrypted_reasoning_signature(encrypted) + assert "thinking_blocks" not in result.model_dump()["choices"][0]["message"] assert message.tool_calls[0].function.name == "get_weather" + anthropic_response: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic( + response=result, tool_name_mapping=None + ) + content: Final = anthropic_response["content"] + assert [block["type"] for block in content] == ["redacted_thinking", "tool_use"] + assert content[0]["data"] == encrypted_reasoning_signature(encrypted) + def test_streaming_function_call_tool_id_for_degenerate_call_id(): """In streaming, Bedrock Mantle's function_call event carries a unique ``id`` diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 49a8341ce9e..bc357878944 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -4,7 +4,9 @@ import json from types import MappingProxyType from typing import Any, Final, cast +import httpx import pytest +import respx import litellm @@ -31,6 +33,7 @@ from litellm.llms.anthropic.pass_through.messages.mid_conversation_system import from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.llms.anthropic import ( AnthopicMessagesAssistantMessageParam, + AnthropicMessagesRequest, AnthropicMessagesUserMessageParam, ) from litellm.types.llms.openai import ChatCompletionAssistantToolCall @@ -647,26 +650,50 @@ def test_translate_openai_reasoning_items_are_not_duplicated_when_thinking_block assert [block["type"] for block in result] == ["redacted_thinking", "text"] -def test_streaming_chunk_with_only_reasoning_items_opens_a_signed_thinking_block(): - encrypted: Final = "gAAAA_delta" - choices: Final = [ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_items=[ - {"type": "reasoning", "id": "rs_1", "encrypted_content": encrypted, "summary": []} - ] +def _replayed_reasoning_request(tools: list[dict[str, object]]) -> AnthropicMessagesRequest: + return AnthropicMessagesRequest( + model="azure/gpt-5.4-nano", + max_tokens=64, + messages=[ + AnthropicMessagesUserMessageParam(role="user", content="Which city is colder?"), + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[ + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_azure_tools")}, + {"type": "tool_use", "id": "call_1", "name": "get_weather", "input": {"city": "Oslo"}}, + ], ), - ) - ] - - block_type, start = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_chunk_to_anthropic_content_block( - choices=choices + AnthropicMessagesUserMessageParam( + role="user", + content=[{"type": "tool_result", "tool_use_id": "call_1", "content": "-7 C"}], + ), + ], + tools=tools, ) - assert block_type == "redacted_thinking" - assert start["data"] == encrypted_reasoning_signature(encrypted) + +@pytest.mark.parametrize( + ("tools", "kept"), + [ + pytest.param( + [{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}], + True, + id="function_tools_bridge_to_responses", + ), + pytest.param([], False, id="no_tools_stay_on_chat"), + ], +) +def test_encrypted_reasoning_is_kept_when_the_request_tools_route_it_to_responses( + tools: list[dict[str, object]], kept: bool +) -> None: + openai_request, _mapping = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + _replayed_reasoning_request(tools), + custom_llm_provider="azure", + api_base="https://keep-check.openai.azure.com", + ) + + assistant: Final = openai_request["messages"][1] + assert bool(assistant.get("thinking_blocks")) is kept def test_translate_anthropic_messages_to_openai_sets_reasoning_content(): @@ -5395,3 +5422,147 @@ def test_eager_input_streaming_tool_reaches_bedrock_converse_as_beta(): assert data["additionalModelRequestFields"]["anthropic_beta"] == ["fine-grained-tool-streaming-2025-05-14"] assert data["toolConfig"]["tools"][0]["toolSpec"]["inputSchema"]["json"] == EAGER_INPUT_SCHEMA + + +@pytest.fixture +def _chat_bridge_httpx(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + +_RETRY_API_BASE: Final = "https://encrypted-retry.openai.azure.com" +_RETRY_URL: Final = r".*encrypted-retry\.openai\.azure\.com/openai/.*responses.*" +_RETRY_TOOLS: Final = [{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}] +_RETRY_REJECTION: Final = httpx.Response( + 400, + json={ + "error": { + "message": "The encrypted content gAAA_foreign could not be verified.", + "type": "invalid_request_error", + "param": None, + "code": "invalid_encrypted_content", + } + }, +) +_UNRELATED_REJECTION: Final = httpx.Response( + 400, json={"error": {"message": "max_output_tokens is too small", "type": "invalid_request_error"}} +) +_RETRY_REPLY: Final = { + "id": "resp_retry", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.4-nano", + "output": [ + { + "type": "message", + "id": "msg_retry", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Oslo", "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 2, "total_tokens": 12}, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], +} + + +def _history_with_foreign_reasoning() -> list[dict[str, object]]: + return [ + {"role": "user", "content": "Which city is colder?"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAA_foreign")}, + {"type": "tool_use", "id": "call_1", "name": "get_weather", "input": {"city": "Oslo"}}, + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "-7 C"}]}, + ] + + +def _sent_reasoning_items(request: httpx.Request) -> list[object]: + sent: Final = json.loads(request.read()) + return [item for item in sent["input"] if item.get("type") == "reasoning"] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_async_messages_retry_once_without_reasoning_the_backend_cannot_decrypt( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=_RETRY_URL).mock( + side_effect=[_RETRY_REJECTION, httpx.Response(200, json=_RETRY_REPLY)] + ) + + response: Final = await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + ) + + assert route.call_count == 2 + assert len(_sent_reasoning_items(route.calls[0].request)) == 1 + assert _sent_reasoning_items(route.calls[1].request) == [] + assert [block["text"] for block in response["content"] if block["type"] == "text"] == ["Oslo"] + + +@pytest.mark.usefixtures("_chat_bridge_httpx") +def test_sync_messages_retry_once_without_reasoning_the_backend_cannot_decrypt( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=_RETRY_URL).mock( + side_effect=[_RETRY_REJECTION, httpx.Response(200, json=_RETRY_REPLY)] + ) + + response: Final = litellm.anthropic.messages.create( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + ) + + assert route.call_count == 2 + assert _sent_reasoning_items(route.calls[1].request) == [] + assert [block["text"] for block in response["content"] if block["type"] == "text"] == ["Oslo"] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +@pytest.mark.parametrize( + ("messages", "rejection"), + [ + pytest.param( + [{"role": "user", "content": "Which city is colder?"}], + _RETRY_REJECTION, + id="no_encrypted_reasoning_to_drop", + ), + pytest.param(_history_with_foreign_reasoning(), _UNRELATED_REJECTION, id="unrelated_rejection"), + ], +) +async def test_messages_do_not_retry_other_rejections( + respx_mock: respx.MockRouter, messages: list[dict[str, object]], rejection: httpx.Response +) -> None: + route: Final = respx_mock.post(url__regex=_RETRY_URL).mock(return_value=rejection) + + with pytest.raises(litellm.BadRequestError): + await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=messages, + ) + + assert route.call_count == 1 From 03fc76ed8d2209a59d6df3b8c7ea44ee6d9aede5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 14:31:42 -0800 Subject: [PATCH 06/41] fix(bedrock): keep skipping a non-string runtime endpoint Backport to lts/1.108.x ahead of its main PR. Cherry-picked from ec312bb5e9 (litellm_fix_bedrock_empty_api_base). --- litellm/llms/bedrock/base_aws_llm.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 2c72a491b3c..0c4151b4b1b 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1526,7 +1526,10 @@ class BaseAWSLLM(SignsRequestsWithAWS): aws_region_name: str, endpoint_type: Literal["runtime", "agent", "agentcore"] | None = "runtime", ) -> tuple[str, str]: - runtime_endpoint: Final = aws_bedrock_runtime_endpoint or get_secret_str("AWS_BEDROCK_RUNTIME_ENDPOINT") + configured_endpoint: Final = ( + aws_bedrock_runtime_endpoint if isinstance(aws_bedrock_runtime_endpoint, str) else None + ) + runtime_endpoint: Final = configured_endpoint or get_secret_str("AWS_BEDROCK_RUNTIME_ENDPOINT") endpoint_url: Final = ( api_base or runtime_endpoint From 6c95a2d970867fa750f04f33126911f491f2a6dd Mon Sep 17 00:00:00 2001 From: yujonglee Date: Sat, 7 Nov 2026 07:44:55 +0900 Subject: [PATCH 07/41] fix(guardrails): check a model-level post_call response once across the proxy and deployment layers (#50300) Backport of #50300 to lts/1.108.x Cherry-picked from 2b4485dee53 (main) Conflicts in litellm/integrations/custom_guardrail.py and litellm/proxy/utils.py came from code the line does not have (the pydantic import, _REQUEST_FIELDS and _guardrails_listed_in in custom_guardrail.py, the stream token-id helpers in proxy/utils.py). Kept the line's code, added the pick's helpers and pydantic import, and defined _REQUEST_FIELDS in both files next to the line's own module-level TypeAdapter constants --- litellm/integrations/custom_guardrail.py | 132 +++- litellm/proxy/utils.py | 8 + .../test_custom_guardrail_post_call_once.py | 566 ++++++++++++++++++ .../proxy_logging/test_guardrail_pipeline.py | 1 + 4 files changed, 690 insertions(+), 17 deletions(-) create mode 100644 tests/unit/integrations/test_custom_guardrail_post_call_once.py diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 1cf71cf59e7..95c8063d7d2 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,6 +1,7 @@ import contextvars import copy import hashlib +import json import os import secrets from collections.abc import Iterable, Mapping, Sequence @@ -9,6 +10,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args import httpx +from pydantic import BaseModel, TypeAdapter from litellm._internal_context import with_service_target from litellm._logging import verbose_logger @@ -71,6 +73,7 @@ GUARDRAIL_SESSIONS_TARGET: Final = "guardrail_sessions" _PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16) _PRE_CALL_RAISED_TOKEN: Final = secrets.token_hex(16) _DEPLOYMENT_HOOK_RAN_TOKEN: Final = secrets.token_hex(16) +_POST_CALL_CHECKED_TOKEN: Final = secrets.token_hex(16) _GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422}) @@ -239,6 +242,98 @@ def _request_carries_marker(data: Mapping[str, object], marker: str | None) -> b return False +def _record_guardrail_marker( + data: dict[str, object], # mutable-ok: the request payload the marker is recorded on + marker: str | None, +) -> None: + if marker is None: + return + for meta_key in ("metadata", "litellm_metadata"): + meta = data.get(meta_key) + if isinstance(meta, dict): + executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) + if isinstance(executed, list): + if marker not in executed: + executed.append(marker) + else: + meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker] + return + data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]} + + +_RESPONSE_METADATA_KEYS: Final = frozenset( + { + "id", + "created", + "created_at", + "model", + "usage", + "usage_info", + "system_fingerprint", + "_hidden_params", + "_response_ms", + } +) + + +def _response_fingerprint(response: object) -> str | None: + """Digest of everything but metadata set between layers, so an unknown field counts as content.""" + try: + fields: Final[Mapping[str, object]] = ( + response.model_dump(mode="json") + if isinstance(response, BaseModel) + else _REQUEST_FIELDS.validate_python(response) + ) + content: Final = {key: value for key, value in fields.items() if key not in _RESPONSE_METADATA_KEYS} + if not content: + return None + return hashlib.sha256( + json.dumps(content, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + ).hexdigest() + except (TypeError, ValueError): + return None + + +_RESPONSE_GUARDRAIL_CHECKS_KEY: Final = "guardrail_checks" + + +def _str_keyed_dict(value: object) -> Mapping[str, object] | None: + if not isinstance(value, dict): + return None + return cast("Mapping[str, object]", value) # cast-ok: response bodies and _hidden_params are keyed by field name + + +def _response_hidden_params(response: object) -> Mapping[str, object] | None: + as_dict: Final = _str_keyed_dict(response) + return _str_keyed_dict( + as_dict.get("_hidden_params") if as_dict is not None else getattr(response, "_hidden_params", None) + ) + + +def _write_response_marker(response: object, marker: str | None) -> None: + """Record a check on the response itself, so it survives the router copying request metadata per hop.""" + if marker is None: + return + hidden: Final = _response_hidden_params(response) + if not isinstance(hidden, dict): + return + hidden[_RESPONSE_GUARDRAIL_CHECKS_KEY] = tuple( + _str_collection(hidden.get(_RESPONSE_GUARDRAIL_CHECKS_KEY)) | {marker} + ) + + +def _response_carries_marker(response: object, marker: str | None) -> bool: + hidden: Final = _response_hidden_params(response) + return ( + marker is not None + and hidden is not None + and marker in _str_collection(hidden.get(_RESPONSE_GUARDRAIL_CHECKS_KEY)) + ) + + +_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object]) + + class CustomGuardrail(CustomLogger): # If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path. use_native_during_call_hook: ClassVar[bool] = False @@ -867,20 +962,7 @@ class CustomGuardrail(CustomLogger): With ``raised``, records instead that the hook raised as a policy pipeline step the pipeline then moved past. The deployment-level hook still runs it. """ - marker: Final = self._pre_call_marker(raised=raised) - if marker is None: - return - for meta_key in ("metadata", "litellm_metadata"): - meta = data.get(meta_key) - if isinstance(meta, dict): - executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) - if isinstance(executed, list): - if marker not in executed: - executed.append(marker) - else: - meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker] - return - data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]} + _record_guardrail_marker(data, self._pre_call_marker(raised=raised)) def pre_call_hook_ran_or_raised(self, data: Mapping[str, object]) -> bool: """Whether this guardrail's ``async_pre_call_hook`` ran for this request, whatever it decided.""" @@ -1077,10 +1159,26 @@ class CustomGuardrail(CustomLogger): cast("dict[str, object]", request_data) # cast-ok: proxy request payloads are str-keyed ) - if not self._is_valid_response_type(result): - return None + returned: Final = result if self._is_valid_response_type(result) else None + if target is self: + return returned + checked: Final = returned if returned is not None else response + marker: Final = self._post_call_check_marker(checked) + _record_guardrail_marker( + cast("dict[str, object]", request_data), # cast-ok: proxy request payloads are str-keyed + marker, + ) + _write_response_marker(checked, marker) + return returned - return result + def _post_call_check_marker(self, response: object) -> str | None: + fingerprint: Final = _response_fingerprint(response) + return self._marker(f"{_POST_CALL_CHECKED_TOKEN}:{fingerprint}") if fingerprint is not None else None + + def post_call_already_checked(self, data: Mapping[str, object], response: object) -> bool: + """Whether a deployment-level hook of this request already ran this guardrail on this exact response.""" + marker: Final = self._post_call_check_marker(response) + return _request_carries_marker(data, marker) or _response_carries_marker(response, marker) async def async_logging_hook( self, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8479f8e45b9..45ca9fb737a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -328,6 +328,7 @@ _VIEW_SETUP_POLL_INTERVAL_SECONDS: Final = 5.0 _VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0 _VIEW_SETUP_GATE_TABLE: Final = '"LiteLLM_SpendLogs"' _VIEW_SETUP_GATE_PROBE_ROWS: Final = TypeAdapter(tuple[Mapping[str, bool], ...]) +_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object]) _ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"] _ViewSetupAttempt: TypeAlias = Literal["ready", "table_missing"] | Exception @@ -4143,6 +4144,7 @@ class ProxyLogging: response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call") + request_view: Final = _REQUEST_FIELDS.validate_python(data) guardrail_callbacks, other_callbacks = _partition_post_call_callbacks() try: # Merge model-level guardrails before checking which guardrails to run @@ -4173,6 +4175,9 @@ class ProxyLogging: ): continue + if callback.post_call_already_checked(request_view, response) is True: + continue + guardrail_response: Any | None = None if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks: @@ -4246,10 +4251,13 @@ class ProxyLogging: immediately before awaiting, and the unified hook pops it before its first suspension point, so concurrent guardrails never race on that key. """ + request_view: Final = _REQUEST_FIELDS.validate_python(data) async def _run_one(callback: CustomGuardrail) -> None: if callback.should_run_guardrail(data=guardrail_data, event_type=GuardrailEventHooks.post_call) is not True: return + if callback.post_call_already_checked(request_view, response) is True: + return if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks: data["guardrail_to_apply"] = callback await self._run_guardrail_with_metrics( diff --git a/tests/unit/integrations/test_custom_guardrail_post_call_once.py b/tests/unit/integrations/test_custom_guardrail_post_call_once.py new file mode 100644 index 00000000000..fb823156929 --- /dev/null +++ b/tests/unit/integrations/test_custom_guardrail_post_call_once.py @@ -0,0 +1,566 @@ +import asyncio +import json +from collections.abc import AsyncIterable, Callable, Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal +from unittest.mock import AsyncMock, patch + +import httpx +import pytest +import respx +from fastapi import HTTPException, Request, Response +from fastapi.responses import StreamingResponse +from mcp.types import Tool as MCPTool +from pydantic import BaseModel, TypeAdapter + +import litellm +from litellm.caching.caching import Cache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import ProxyLogging +from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler +from litellm.router import Router +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CallTypes, Choices, GenericGuardrailAPIInputs, LLMResponseTypes, Message, ModelResponse + +pytestmark = pytest.mark.respx(assert_all_called=False) + +REPLY: Final = "hi from the model" +REWRITTEN: Final = "rewritten by the proxy" +MASKED: Final = "masked by G" + +Route = Literal["acompletion", "aresponses", "anthropic_messages"] + +_OPENAI_API_BASE: Final = "https://api.openai.com/v1" +_FAILING_API_BASE: Final = "https://primary.invalid/v1" +_MCP_API_BASE: Final = "https://mcp-rounds.invalid/v1" + +_CHAT_REPLY: Final[Mapping[str, object]] = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 0, + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": REPLY}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +_CHAT_SSE: Final = ( + "".join( + f"data: {json.dumps(chunk)}\n\n" + for chunk in ( + { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 0, + "model": "gpt-4o", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": REPLY}, "finish_reason": None}], + }, + { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 0, + "model": "gpt-4o", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, + ) + ) + + "data: [DONE]\n\n" +) +_RESPONSES_REPLY: Final[Mapping[str, object]] = { + "id": "resp_1", + "object": "response", + "created_at": 0, + "status": "completed", + "model": "gpt-4o", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": REPLY, "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, +} +_RESPONSES_SSE: Final = "".join( + f"data: {json.dumps(event)}\n\n" + for event in ( + {"type": "response.created", "sequence_number": 0, "response": {**_RESPONSES_REPLY, "output": []}}, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": REPLY, + }, + {"type": "response.completed", "sequence_number": 2, "response": _RESPONSES_REPLY}, + ) +) +_TOOL_CALL_ROUND: Final[Mapping[str, object]] = { + **_RESPONSES_REPLY, + "id": "resp_tool_call", + "output": [ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_echo", + "name": "echo", + "arguments": '{"value": "PING"}', + "status": "completed", + } + ], +} +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) + + +class _CountedGuardrail(CustomGuardrail): + checked: tuple[str, ...] = () + + +class ApplyGuardrailG(_CountedGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + self.checked = (*self.checked, " ".join(inputs.get("texts") or ())) + return inputs + + +def _reply_text(response: LLMResponseTypes) -> str: + match response: + case ModelResponse(choices=[Choices(message=Message(content=str() as text))]): + return text + case _: + return str(response) + + +class OwnHookG(_CountedGuardrail): + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + ) -> None: + self.checked = (*self.checked, _reply_text(response)) + + +class InheritedApplyGuardrailG(ApplyGuardrailG): + pass + + +class MaskingG(ApplyGuardrailG): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + checked: Final = await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + return {**checked, "texts": [MASKED for _ in checked.get("texts") or ()]} + + +class BlockingG(ApplyGuardrailG): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + raise HTTPException(status_code=400, detail="blocked by G") + + +class RewritingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "texts": [REWRITTEN for _ in inputs.get("texts") or ()]} + + +class ToolCallBlockingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + if input_type == "response" and inputs.get("tool_calls"): + raise HTTPException(status_code=400, detail="tool calls are not allowed") + return inputs + + +def _guardrail_g( + style: type[_CountedGuardrail] = ApplyGuardrailG, default_on: bool = False, run_in_parallel: bool = False +) -> _CountedGuardrail: + return style( + guardrail_name="G", + event_hook=GuardrailEventHooks.post_call, + default_on=default_on, + run_in_parallel=run_in_parallel, + ) + + +def _known_gap(action_item: str) -> pytest.MarkDecorator: + return pytest.mark.xfail( + raises=AssertionError, + strict=True, + reason=f"broken before #40154, out of scope for LIT-10481, tracked as RCA {action_item}", + ) + + +def _deployment(model_name: str, guardrails: tuple[str, ...], api_base: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-test", + "api_base": api_base, + **({"guardrails": list(guardrails)} if guardrails else {}), + }, + } + + +def _router(primary: tuple[str, ...], fallback: tuple[str, ...] | None = None, cached: bool = False) -> Router: + if fallback is None: + return Router(model_list=[_deployment("gpt", primary, _OPENAI_API_BASE)], cache_responses=cached) + return Router( + model_list=[ + _deployment("gpt", primary, _FAILING_API_BASE), + _deployment("gpt-backup", fallback, _OPENAI_API_BASE), + ], + fallbacks=[{"gpt": ["gpt-backup"]}], + num_retries=0, + ) + + +def _json_or_sse(json_reply: Mapping[str, object], sse_reply: str) -> Callable[[httpx.Request], httpx.Response]: + def reply(request: httpx.Request) -> httpx.Response: + if _JSON_OBJECT.validate_json(request.content).get("stream"): + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=sse_reply.encode()) + return httpx.Response(200, json=dict(json_reply)) + + return reply + + +async def _drain(stream: AsyncIterable[object]) -> str: + return "".join([str(chunk) async for chunk in stream]) + + +@pytest.fixture(autouse=True) +def provider(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> Iterator[respx.MockRouter]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + ProxyLogging._callback_capabilities_cache.clear() # pyright: ignore[reportPrivateUsage] # keyed on id(), which a collected guardrail from an earlier test can reuse + for endpoint in ("chat/completions", "responses"): + respx_mock.post(f"{_FAILING_API_BASE}/{endpoint}").mock( + return_value=httpx.Response(500, json={"error": {"message": "primary down"}}) + ) + respx_mock.post(f"{_OPENAI_API_BASE}/chat/completions").mock(side_effect=_json_or_sse(_CHAT_REPLY, _CHAT_SSE)) + respx_mock.post(f"{_OPENAI_API_BASE}/responses").mock(side_effect=_json_or_sse(_RESPONSES_REPLY, _RESPONSES_SSE)) + yield respx_mock + litellm.in_memory_llm_clients_cache.flush_cache() + + +@dataclass(frozen=True, slots=True) +class ProxyCase: + route: Route = "acompletion" + stream: bool = False + primary: tuple[str, ...] = ("G",) + fallback: tuple[str, ...] | None = None + request: tuple[str, ...] = () + default_on: bool = False + style: type[_CountedGuardrail] = ApplyGuardrailG + cached: bool = False + client_sees: str = REPLY + checks: tuple[str, ...] = (REPLY,) + + +PROXY_CASES: Final = ( + pytest.param(ProxyCase(), id="A-chat"), + pytest.param(ProxyCase(route="aresponses"), id="A-responses"), + pytest.param(ProxyCase(route="anthropic_messages"), id="A-messages"), + pytest.param(ProxyCase(request=("G",)), id="A-chat-model-and-request"), + pytest.param(ProxyCase(style=MaskingG, client_sees=MASKED), id="A-chat-g-masks-its-own-output"), + pytest.param(ProxyCase(stream=True), id="B-chat-stream"), + pytest.param(ProxyCase(route="aresponses", stream=True), id="B-responses-stream"), + pytest.param(ProxyCase(route="anthropic_messages", stream=True), id="B-messages-stream"), + pytest.param(ProxyCase(primary=(), fallback=("G",)), id="E-only-on-fallback"), + pytest.param( + ProxyCase(request=("rewriter",), client_sees=REWRITTEN, checks=(REPLY, REWRITTEN)), + id="G-rewritten-after-check", + ), + pytest.param(ProxyCase(fallback=("G",)), id="H-on-both-deployments"), + pytest.param(ProxyCase(primary=(), request=("G",)), id="I-request-chat"), + pytest.param(ProxyCase(route="aresponses", primary=(), request=("G",)), id="I-request-responses"), + pytest.param(ProxyCase(style=OwnHookG), id="J-A-chat-own-hook", marks=_known_gap("AI-5")), + pytest.param(ProxyCase(primary=(), fallback=("G",), style=OwnHookG), id="E-only-on-fallback-own-hook"), + pytest.param(ProxyCase(primary=(), request=("G",), style=OwnHookG), id="I-request-chat-own-hook"), + pytest.param( + ProxyCase(primary=(), request=("G",), style=InheritedApplyGuardrailG), + id="K-I-request-chat-inherited-apply-guardrail", + marks=_known_gap("AI-2"), + ), + pytest.param(ProxyCase(primary=("other",), default_on=True), id="L-default-on"), +) + +_ROUTE_PATHS: Final[Mapping[Route, str]] = MappingProxyType( + { + "acompletion": "/v1/chat/completions", + "aresponses": "/v1/responses", + "anthropic_messages": "/v1/messages", + } +) +_ROUTE_BODIES: Final[Mapping[Route, Mapping[str, object]]] = MappingProxyType( + { + "acompletion": {"messages": [{"role": "user", "content": "hello"}]}, + "aresponses": {"input": "hello"}, + "anthropic_messages": {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 32}, + } +) + + +async def _proxy_reply(case: ProxyCase) -> str: + from litellm.proxy.proxy_server import ProxyConfig, select_data_generator + + path: Final = _ROUTE_PATHS[case.route] + router: Final = _router(case.primary, case.fallback, case.cached) + response: Final = await ProxyBaseLLMRequestProcessing( + data={ + "model": "gpt", + **_ROUTE_BODIES[case.route], + **({"stream": True} if case.stream else {}), + **({"guardrails": list(case.request)} if case.request else {}), + } + ).base_process_llm_request( + request=Request({"type": "http", "method": "POST", "path": path, "headers": [], "query_string": b""}), + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", request_route=path), + route_type=case.route, + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + general_settings={}, + proxy_config=ProxyConfig(), + select_data_generator=select_data_generator, + llm_router=router, + is_streaming_request=case.stream, + ) + text: Final = ( + await _drain(response.body_iterator) + if isinstance(response, StreamingResponse) + else json.dumps(response.model_dump() if isinstance(response, BaseModel) else response) + ) + await asyncio.sleep(0) + return text + + +@pytest.mark.parametrize("case", PROXY_CASES) +async def test_a_proxy_request_checks_each_distinct_response_once( + monkeypatch: pytest.MonkeyPatch, case: ProxyCase +) -> None: + guardrail: Final = _guardrail_g(case.style, default_on=case.default_on) + rewriter: Final = RewritingGuardrail(guardrail_name="rewriter", event_hook=GuardrailEventHooks.post_call) + monkeypatch.setattr(litellm, "callbacks", [rewriter, guardrail]) + + reply: Final = await _proxy_reply(case) + + assert case.client_sees in reply + assert sorted(guardrail.checked) == sorted(case.checks) + + +@pytest.mark.parametrize("route", ["acompletion", "aresponses"]) +async def test_a_cache_hit_is_checked_like_a_fresh_response( + monkeypatch: pytest.MonkeyPatch, provider: respx.MockRouter, route: Route +) -> None: + guardrail: Final = _guardrail_g() + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr(litellm, "cache", Cache(type="local")) + + case: Final = ProxyCase(route=route, cached=True) + replies: Final = (await _proxy_reply(case), await _proxy_reply(case)) + + assert provider.calls.call_count == 1, "the second request must be served from the cache" + assert all(REPLY in reply for reply in replies) + assert guardrail.checked == (REPLY, REPLY) + + +async def test_a_blocking_verdict_blocks_the_proxy_response_after_one_check(monkeypatch: pytest.MonkeyPatch) -> None: + guardrail: Final = _guardrail_g(BlockingG) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + with pytest.raises(HTTPException, check=lambda e: e.detail == "blocked by G"): + await _proxy_reply(ProxyCase()) + + assert guardrail.checked == (REPLY,) + + +SdkCaller = Literal["sdk-call", "sdk-router"] + + +@dataclass(frozen=True, slots=True) +class SdkCase: + caller: SdkCaller = "sdk-call" + stream: bool = False + style: type[_CountedGuardrail] = ApplyGuardrailG + + +SDK_CASES: Final = ( + pytest.param(SdkCase(), id="C-sdk-call"), + pytest.param(SdkCase(caller="sdk-router"), id="C-sdk-router"), + pytest.param(SdkCase(style=OwnHookG), id="C-sdk-call-own-hook"), + pytest.param(SdkCase(stream=True), id="D-sdk-call-stream", marks=_known_gap("AI-4")), +) + + +async def _sdk_reply(case: SdkCase) -> str: + messages: Final = [{"role": "user", "content": "hello"}] + response: Final = await ( + _router(("G",)).acompletion(model="gpt", messages=messages, stream=case.stream) + if case.caller == "sdk-router" + else litellm.acompletion( + model="openai/gpt-4o", api_key="sk-test", guardrails=["G"], messages=messages, stream=case.stream + ) + ) + text: Final = await _drain(response) if isinstance(response, AsyncIterable) else str(response) + await asyncio.sleep(0) + return text + + +@pytest.mark.parametrize("case", SDK_CASES) +async def test_an_sdk_call_checks_its_response_once(monkeypatch: pytest.MonkeyPatch, case: SdkCase) -> None: + guardrail: Final = _guardrail_g(case.style) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + reply: Final = await _sdk_reply(case) + + assert REPLY in reply + assert guardrail.checked == (REPLY,) + + +async def test_mcp_auto_execute_checks_the_tool_call_round_before_running_its_tools( + monkeypatch: pytest.MonkeyPatch, provider: respx.MockRouter +) -> None: + monkeypatch.setattr( + litellm, "callbacks", [ToolCallBlockingGuardrail(guardrail_name="G", event_hook=GuardrailEventHooks.post_call)] + ) + provider.post(f"{_MCP_API_BASE}/responses").mock( + side_effect=[httpx.Response(200, json=_TOOL_CALL_ROUND), httpx.Response(200, json=_RESPONSES_REPLY)] + ) + echo_tool: Final = MCPTool.model_validate( + {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, + by_name=False, + ) + run_tools: Final = AsyncMock(return_value=[{"tool_call_id": "call_echo", "result": "PING", "name": "echo"}]) + + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_without_openai_transform", + AsyncMock(return_value=([echo_tool], {"echo": "demo"})), + ), + patch.object(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", run_tools), + pytest.raises(HTTPException, check=lambda e: e.detail == "tool calls are not allowed"), + ): + await litellm.aresponses( + model="openai/gpt-4o", + input="Call echo with PING", + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + api_key="sk-test", + api_base=_MCP_API_BASE, + guardrails=["G"], + proxy_server_request={"url": "http://localhost:4000/v1/responses", "method": "POST"}, + litellm_metadata={"user_api_key_request_route": "/v1/responses"}, + ) + + run_tools.assert_not_awaited() + + +def _same_response(reply: ModelResponse) -> ModelResponse: + return reply + + +def _model_renamed(reply: ModelResponse) -> ModelResponse: + return reply.model_copy(update={"model": "public-alias"}) + + +def _rewritten(_: ModelResponse) -> ModelResponse: + return ModelResponse(choices=[{"message": {"role": "assistant", "content": REWRITTEN}}]) + + +PROXY_LOOP_CASES: Final = ( + pytest.param(_same_response, (REPLY,), id="A-same-response"), + pytest.param(_model_renamed, (REPLY,), id="A-model-renamed"), + pytest.param(_rewritten, (REPLY, REWRITTEN), id="G-rewritten"), +) + + +@pytest.mark.parametrize("parallel", [False, True], ids=["sequential", "parallel"]) +@pytest.mark.parametrize(("proxy_sees", "checks"), PROXY_LOOP_CASES) +async def test_each_proxy_post_call_loop_skips_only_output_the_deployment_checked( + monkeypatch: pytest.MonkeyPatch, + parallel: bool, + proxy_sees: Callable[[ModelResponse], ModelResponse], + checks: tuple[str, ...], +) -> None: + guardrail: Final = _guardrail_g(run_in_parallel=parallel) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + data: Final = { + "model": "gpt", + "proxy_server_request": {"url": "http://localhost:4000/v1/chat/completions"}, + "metadata": {"guardrails": ["G"]}, + } + reply: Final = ModelResponse(model="provider-model", choices=[{"message": {"role": "assistant", "content": REPLY}}]) + + await guardrail.async_post_call_success_deployment_hook( + request_data={**data, "guardrails": ["G"]}, response=reply, call_type=CallTypes.acompletion + ) + await ProxyLogging(user_api_key_cache=UserApiKeyCache()).post_call_success_hook( + data=data, response=proxy_sees(reply), user_api_key_dict=UserAPIKeyAuth() + ) + + assert sorted(guardrail.checked) == sorted(checks) + + +def _ocr_reply(markdown: str) -> OCRResponse: + return OCRResponse.model_validate( + {"pages": [{"index": 0, "markdown": markdown, "images": [], "dimensions": None}], "model": "mistral-ocr-latest"} + ) + + +@pytest.mark.parametrize( + ("proxy_sees", "checks"), + [ + pytest.param(REPLY, (REPLY,), id="A-ocr-same-pages"), + pytest.param(REWRITTEN, (REPLY, REWRITTEN), id="G-ocr-rewritten-pages"), + ], +) +async def test_an_ocr_response_is_skipped_only_when_its_pages_were_checked( + monkeypatch: pytest.MonkeyPatch, proxy_sees: str, checks: tuple[str, ...] +) -> None: + guardrail: Final = _guardrail_g() + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + data: Final = { + "model": "ocr", + "document": {"type": "document_url", "document_url": "https://example.com/doc.pdf"}, + "proxy_server_request": {"url": "http://localhost:4000/v1/ocr"}, + "metadata": {"guardrails": ["G"], "user_api_key_request_route": "/v1/ocr"}, + } + + await guardrail.async_post_call_success_deployment_hook( + request_data={**data, "guardrails": ["G"]}, response=_ocr_reply(REPLY), call_type=CallTypes.aocr + ) + await ProxyLogging(user_api_key_cache=UserApiKeyCache()).post_call_success_hook( + data=data, response=_ocr_reply(proxy_sees), user_api_key_dict=UserAPIKeyAuth(request_route="/v1/ocr") + ) + + assert sorted(guardrail.checked) == sorted(checks) diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index fb3060b42a9..8d106531218 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -696,6 +696,7 @@ def _moderation_guardrail() -> MagicMock: cb.should_run_guardrail = MagicMock(return_value=True) cb.async_moderation_hook = AsyncMock(return_value=None) cb.async_post_call_success_hook = AsyncMock(return_value=None) + cb.post_call_already_checked = MagicMock(return_value=False) cb.run_in_parallel = False return cb From 1fa046752a1df092b9cac7b9c9bc262556a334ca Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 15:02:52 -0800 Subject: [PATCH 08/41] fix(anthropic): let Python 3.10 import the backported messages handler and clear its new lint and type errors Quote the rejected-reasoning return annotation, since a string TypeAlias cannot take `| None` at definition time before Python 3.14. Return only the response from the retry helpers, since this line's stream wrapper takes no completion request. Type the bridge check result and the content block union the way main does, and validate the object tuple through TypeAdapter. --- .../prompt_templates/common_utils.py | 4 +++- .../pass_through/adapters/handler.py | 20 +++++++++---------- litellm/responses/bridge_check.py | 2 +- litellm/types/llms/anthropic.py | 4 +++- 4 files changed, 16 insertions(+), 14 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index f5fa6f649a6..f9c4ebca350 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -2292,11 +2292,13 @@ def encrypted_content_of_block(block: Mapping[str, object]) -> str | None: _REASONING_SUMMARY_SEPARATOR: Final = "\n\n" +_OBJECT_TUPLE: Final = TypeAdapter(tuple[object, ...]) + def _object_sequence(value: object) -> tuple[object, ...]: if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): return () - return tuple(value) + return _OBJECT_TUPLE.validate_python(value) def _iter_mappings(value: object) -> tuple[Mapping[str, object], ...]: diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index fda9fe61060..e385ee194b3 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -86,7 +86,7 @@ class _ChatTargetReasoning: def _messages_without_rejected_reasoning( messages: _AnthropicMessages, error: litellm.BadRequestError, -) -> _AnthropicMessages | None: +) -> "_AnthropicMessages | None": error_message: Final = str(error) rejected: Final = "invalid_encrypted_content" in error_message or "could not be verified" in error_message if not rejected or not messages_carry_encrypted_reasoning(messages): @@ -99,30 +99,28 @@ async def _acompletion_without_rejected_reasoning( completion_kwargs: _CompletionKwargs, messages: _AnthropicMessages, prepare_replay: _PrepareReplay, -) -> tuple[_CompletionKwargs, "ChatResult"]: +) -> "ChatResult": try: - return completion_kwargs, await litellm.acompletion(**completion_kwargs) + return await litellm.acompletion(**completion_kwargs) except litellm.BadRequestError as error: replayed_messages: Final = _messages_without_rejected_reasoning(messages, error) if replayed_messages is None: raise - replay_kwargs: Final = prepare_replay(replayed_messages) - return replay_kwargs, await litellm.acompletion(**replay_kwargs) + return await litellm.acompletion(**prepare_replay(replayed_messages)) def _completion_without_rejected_reasoning( completion_kwargs: _CompletionKwargs, messages: _AnthropicMessages, prepare_replay: _PrepareReplay, -) -> tuple[_CompletionKwargs, "ChatResult | Coroutine[object, object, ChatResult]"]: +) -> "ChatResult | Coroutine[object, object, ChatResult]": try: - return completion_kwargs, litellm.completion(**completion_kwargs) + return litellm.completion(**completion_kwargs) except litellm.BadRequestError as error: replayed_messages: Final = _messages_without_rejected_reasoning(messages, error) if replayed_messages is None: raise - replay_kwargs: Final = prepare_replay(replayed_messages) - return replay_kwargs, litellm.completion(**replay_kwargs) + return litellm.completion(**prepare_replay(replayed_messages)) def _is_compaction_block(block: object) -> bool: @@ -850,7 +848,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: extra_kwargs=kwargs, ) initial_kwargs, tool_name_mapping = prepare(messages=effective_messages) - completion_kwargs, completion_response = await _acompletion_without_rejected_reasoning( + completion_response: Final = await _acompletion_without_rejected_reasoning( initial_kwargs, effective_messages, lambda replayed_messages: prepare(messages=replayed_messages)[0] ) logging_obj: Final = litellm_logging_obj_from_kwargs(kwargs) @@ -994,7 +992,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: extra_kwargs=kwargs, ) initial_kwargs, tool_name_mapping = prepare(messages=effective_messages) - completion_kwargs, completion_response = _completion_without_rejected_reasoning( + completion_response: Final = _completion_without_rejected_reasoning( initial_kwargs, effective_messages, lambda replayed_messages: prepare(messages=replayed_messages)[0] ) logging_obj: Final = litellm_logging_obj_from_kwargs(kwargs) diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index cdd10d489a1..d8a9440bc57 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -72,7 +72,7 @@ def responses_api_bridge_check( reasoning_effort: str | Mapping[str, object] | None = None, reasoning_summary: object | None = None, api_base: str | None = None, -) -> tuple[dict, str]: +) -> tuple[dict[str, object], str]: model_info: dict[str, object] = {} # Global flag: route ALL OpenAI chat completions through Responses API. diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 53abac75162..2f59c1765a5 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -545,7 +545,9 @@ class ContentBlockStartText(TypedDict): content_block: TextBlock -ContentBlockContentBlockDict = ToolUseBlock | TextBlock | ChatCompletionThinkingBlock +ContentBlockContentBlockDict = ( + ToolUseBlock | TextBlock | ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock +) ContentBlockStart = ContentBlockStartToolUse | ContentBlockStartText From 47931a3bf2321a1e9015803d5f988da2a2f2d40c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 7 Nov 2026 07:57:20 +0900 Subject: [PATCH 09/41] chore(deps): lock filelock 4.2.4 on 3.11+ and 3.20.3 on 3.10, as main does Widens the filelock range to <5 and adds main's two filelock constraints, matching main's pyproject.toml since #49204 --- pyproject.toml | 4 +- uv.lock | 432 ++++++++++++++++++++++++++----------------------- 2 files changed, 230 insertions(+), 206 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index fca9ea6c881..ca25b5bc48e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ dependencies = [ # When changing a floor, verify it installs + imports on every supported # Python with: `uv pip install --resolution=lowest-direct .` "fastuuid>=0.14.0,<1.0", - "filelock>=3.16.1,<4.0", + "filelock>=3.16.1,<5", "httpx[http2]>=0.28.0,<1.0", "openai>=2.20.0,<4.0.0", "python-dotenv>=1.0.0,<2.0", @@ -358,6 +358,8 @@ constraint-dependencies = [ "httplib2>=0.32.0", "setuptools>=84.1.0", "dnspython>=2.9.0; python_version >= '3.11'", + "filelock<3.20.4; python_version < '3.11'", + "filelock>=4.2.1; python_version >= '3.11'", ] override-dependencies = [ # a2a-sdk 1.x requires packaging>=24.0; lunary 1.4.x still caps at <24.0. diff --git a/uv.lock b/uv.lock index 0560d07f501..829abd54c5d 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 5 +revision = 3 requires-python = ">=3.10, <3.15" resolution-markers = [ "python_full_version >= '3.14'", @@ -10,17 +10,17 @@ resolution-markers = [ ] [options] -exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. +exclude-newer = "2026-11-03T22:57:01.748805Z" exclude-newer-span = "P3D" [options.exclude-newer-package] -boto3 = "2026-10-21T19:31:00Z" -botocore = "2026-10-21T19:31:00Z" +mlflow-skinny = "2026-10-21T03:00:00Z" brotli = "2026-10-22T14:01:00Z" certifi = "2026-10-20T03:00:00Z" grpcio = "2026-10-22T15:00:00Z" +boto3 = "2026-10-21T19:31:00Z" +botocore = "2026-10-21T19:31:00Z" mlflow = "2026-10-21T03:00:00Z" -mlflow-skinny = "2026-10-21T03:00:00Z" mlflow-tracing = "2026-10-21T03:00:00Z" pypdf = "2026-10-28T17:01:00Z" @@ -33,6 +33,8 @@ members = [ constraints = [ { name = "aiohttp", specifier = ">=3.14.2,<4.0" }, { name = "dnspython", marker = "python_full_version >= '3.11'", specifier = ">=2.9.0" }, + { name = "filelock", marker = "python_full_version < '3.11'", specifier = "<3.20.4" }, + { name = "filelock", marker = "python_full_version >= '3.11'", specifier = ">=4.2.1" }, { name = "httplib2", specifier = ">=0.32.0" }, { name = "packaging", specifier = ">=24.0" }, { name = "setuptools", specifier = ">=84.1.0" }, @@ -239,9 +241,9 @@ name = "aiologic" version = "0.17.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "sniffio" }, - { name = "typing-extensions" }, - { name = "wrapt" }, + { name = "sniffio", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "wrapt", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/53/a7/809482759f40079f4c4328c7318bf569ae25d457f5017aad30a1b9aafedc/aiologic-0.17.0.tar.gz", hash = "sha256:65aa058e858c94cd208badb188e7f00b54dcabb3ba85b34f794db98074d108b9", size = 251625, upload-time = "2026-06-14T12:24:35.367Z" } wheels = [ @@ -533,14 +535,14 @@ name = "aurelio-sdk" version = "0.0.19" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "aiofiles" }, - { name = "aiohttp" }, - { name = "colorlog" }, - { name = "pydantic" }, - { name = "python-dotenv" }, - { name = "requests" }, - { name = "requests-toolbelt" }, - { name = "tornado", version = "6.5.10", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.14'" }, + { name = "aiofiles", marker = "python_full_version < '3.14'" }, + { name = "aiohttp", marker = "python_full_version < '3.14'" }, + { name = "colorlog", marker = "python_full_version < '3.14'" }, + { name = "pydantic", marker = "python_full_version < '3.14'" }, + { name = "python-dotenv", marker = "python_full_version < '3.14'" }, + { name = "requests", marker = "python_full_version < '3.14'" }, + { name = "requests-toolbelt", marker = "python_full_version < '3.14'" }, + { name = "tornado", version = "6.5.10", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "tornado", version = "6.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/27/0e/c2e369ad173fb3d76448e46d10beb3dcc53388318933ddf8169a3f21a810/aurelio_sdk-0.0.19.tar.gz", hash = "sha256:14107e7440ff2efd0b4a08c52fb595e7680bd4bc973a0ddfb3b64157c6666b91", size = 15258, upload-time = "2025-03-24T14:37:32.203Z" } @@ -553,9 +555,9 @@ name = "aws-sdk-bedrock-runtime" version = "0.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "smithy-aws-core", extra = ["eventstream", "json"] }, - { name = "smithy-core" }, - { name = "smithy-http", extra = ["aiohttp"] }, + { name = "smithy-aws-core", extra = ["eventstream", "json"], marker = "python_full_version >= '3.12'" }, + { name = "smithy-core", marker = "python_full_version >= '3.12'" }, + { name = "smithy-http", extra = ["aiohttp"], marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/8e/b3/9c225cbfe9f17ea2e3d75a0fdd0b325ef79839b9c09a376bda63a7bf3bb3/aws_sdk_bedrock_runtime-0.11.0.tar.gz", hash = "sha256:f2c45d34625bf6a7b56375e29a53a16b376880bda771e4bbf7d84491622eb193", size = 173854, upload-time = "2026-08-24T21:17:16.304Z" } wheels = [ @@ -564,7 +566,7 @@ wheels = [ [package.optional-dependencies] awscrt = [ - { name = "smithy-http", extra = ["awscrt"] }, + { name = "smithy-http", extra = ["awscrt"], marker = "python_full_version >= '3.12'" }, ] [[package]] @@ -1247,7 +1249,7 @@ name = "colorlog" version = "6.10.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "colorama", marker = "python_full_version < '3.14' and sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a2/61/f083b5ac52e505dfc1c624eafbf8c7589a0d7f32daa398d2e7590efa5fda/colorlog-6.10.1.tar.gz", hash = "sha256:eb4ae5cb65fe7fec7773c2306061a8e63e02efc2c72eba9d27b0fa23c94f1321", size = 17162, upload-time = "2025-10-16T16:14:11.978Z" } wheels = [ @@ -1271,7 +1273,7 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/66/54/eb9bfc647b19f2009dd5c7f5ec51c4e6ca831725f1aea7a993034f483147/contourpy-1.3.2.tar.gz", hash = "sha256:b6945942715a034c671b7fc54f9588126b0b8bf23db2696e3ca8328f3ff0ab54", size = 13466130, upload-time = "2025-04-15T17:47:53.79Z" } wheels = [ @@ -1344,7 +1346,7 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/58/01/1253e6698a07380cd31a736d248a3f2a50a7c88779a1813da27503cadc2a/contourpy-1.3.3.tar.gz", hash = "sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880", size = 13466174, upload-time = "2025-07-26T12:03:12.549Z" } @@ -1614,8 +1616,8 @@ name = "culsans" version = "0.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "aiologic" }, - { name = "typing-extensions" }, + { name = "aiologic", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/d9/e3/49afa1bc180e0d28008ec6bcdf82a4072d1c7a41032b5b759b60814ca4b0/culsans-0.11.0.tar.gz", hash = "sha256:0b43d0d05dce6106293d114c86e3fb4bfc63088cfe8ff08ed3fe36891447fe33", size = 107546, upload-time = "2025-12-31T23:15:38.196Z" } wheels = [ @@ -1898,7 +1900,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -2058,11 +2060,29 @@ wheels = [ [[package]] name = "filelock" -version = "3.32.7" +version = "3.20.3" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/0f/59/e19834834cb01a32febfbb0f8a23a9088088f5d45991824ff2bc3b5e8acb/filelock-3.32.7.tar.gz", hash = "sha256:37b8a3d9811b0f9aef7e5ec5c71bb320de52df51e6ca9bcd6f5ad81187660da7", size = 225154, upload-time = "2026-09-16T00:24:20.907Z" } +resolution-markers = [ + "python_full_version < '3.11'", +] +sdist = { url = "https://files.pythonhosted.org/packages/1d/65/ce7f1b70157833bf3cb851b556a37d4547ceafc158aa9b34b36782f23696/filelock-3.20.3.tar.gz", hash = "sha256:18c57ee915c7ec61cff0ecf7f0f869936c7c30191bb0cf406f1341778d0834e1", size = 19485, upload-time = "2026-01-09T17:55:05.421Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/15/df/31098c5aeb4d966b553641472bd55fcf5fdfac953549894b8a765ba44e91/filelock-3.32.7-py3-none-any.whl", hash = "sha256:65ff0d0190ea42038b32bda4b77834fb05be2cad4c5b9b01aa4dfb3614536e52", size = 100157, upload-time = "2026-09-16T00:24:19.543Z" }, + { url = "https://files.pythonhosted.org/packages/b5/36/7fb70f04bf00bc646cd5bb45aa9eddb15e19437a28b8fb2b4a5249fac770/filelock-3.20.3-py3-none-any.whl", hash = "sha256:4b0dda527ee31078689fc205ec4f1c1bf7d56cf88b6dc9426c4f230e46c2dce1", size = 16701, upload-time = "2026-01-09T17:55:04.334Z" }, +] + +[[package]] +name = "filelock" +version = "4.2.4" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14'", + "python_full_version == '3.13.*'", + "python_full_version == '3.12.*'", + "python_full_version == '3.11.*'", +] +sdist = { url = "https://files.pythonhosted.org/packages/d8/28/3893324bcadf8f89b1616637468d766af3f625008c0735bf704fdfa86617/filelock-4.2.4.tar.gz", hash = "sha256:c6d8663535e1650ee430507b4ffbac7b464dd64b7647ea0112a9f406499d163d", size = 652444, upload-time = "2026-10-19T01:57:51.001Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/1b/5e74289a1878e666bb4b65fe5805a04cafcac2aa5f519ef1ead67ef91c56/filelock-4.2.4-py3-none-any.whl", hash = "sha256:301241438c57936ef93f4243a94f2559022268bfcf9b375a9a72eb0f8db8358e", size = 162475, upload-time = "2026-10-19T01:57:49.419Z" }, ] [[package]] @@ -2481,11 +2501,11 @@ resolution-markers = [ "python_full_version >= '3.14'", ] dependencies = [ - { name = "google-auth" }, - { name = "googleapis-common-protos" }, - { name = "proto-plus" }, - { name = "protobuf" }, - { name = "requests" }, + { name = "google-auth", marker = "python_full_version >= '3.14'" }, + { name = "googleapis-common-protos", marker = "python_full_version >= '3.14'" }, + { name = "proto-plus", marker = "python_full_version >= '3.14'" }, + { name = "protobuf", marker = "python_full_version >= '3.14'" }, + { name = "requests", marker = "python_full_version >= '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/09/cd/63f1557235c2440fe0577acdbc32577c5c002684c58c7f4d770a92366a24/google_api_core-2.25.2.tar.gz", hash = "sha256:1c63aa6af0d0d5e37966f157a77f9396d820fba59f9e43e9415bc3dc5baff300", size = 166266, upload-time = "2025-10-03T00:07:34.778Z" } wheels = [ @@ -2494,8 +2514,8 @@ wheels = [ [package.optional-dependencies] grpc = [ - { name = "grpcio" }, - { name = "grpcio-status" }, + { name = "grpcio", marker = "python_full_version >= '3.14'" }, + { name = "grpcio-status", marker = "python_full_version >= '3.14'" }, ] [[package]] @@ -2509,11 +2529,11 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "google-auth" }, - { name = "googleapis-common-protos" }, - { name = "proto-plus" }, - { name = "protobuf" }, - { name = "requests" }, + { name = "google-auth", marker = "python_full_version < '3.14'" }, + { name = "googleapis-common-protos", marker = "python_full_version < '3.14'" }, + { name = "proto-plus", marker = "python_full_version < '3.14'" }, + { name = "protobuf", marker = "python_full_version < '3.14'" }, + { name = "requests", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/16/ce/502a57fb0ec752026d24df1280b162294b22a0afb98a326084f9a979138b/google_api_core-2.30.3.tar.gz", hash = "sha256:e601a37f148585319b26db36e219df68c5d07b6382cff2d580e83404e44d641b", size = 177001, upload-time = "2026-04-10T00:41:28.035Z" } wheels = [ @@ -2522,8 +2542,8 @@ wheels = [ [package.optional-dependencies] grpc = [ - { name = "grpcio" }, - { name = "grpcio-status" }, + { name = "grpcio", marker = "python_full_version < '3.14'" }, + { name = "grpcio-status", marker = "python_full_version < '3.14'" }, ] [[package]] @@ -2709,12 +2729,12 @@ resolution-markers = [ "python_full_version >= '3.14'", ] dependencies = [ - { name = "google-api-core", version = "2.25.2", source = { registry = "https://pypi.org/simple" } }, - { name = "google-auth" }, - { name = "google-cloud-core" }, - { name = "google-crc32c" }, - { name = "google-resumable-media" }, - { name = "requests" }, + { name = "google-api-core", version = "2.25.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, + { name = "google-auth", marker = "python_full_version >= '3.14'" }, + { name = "google-cloud-core", marker = "python_full_version >= '3.14'" }, + { name = "google-crc32c", marker = "python_full_version >= '3.14'" }, + { name = "google-resumable-media", marker = "python_full_version >= '3.14'" }, + { name = "requests", marker = "python_full_version >= '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/bd/ef/7cefdca67a6c8b3af0ec38612f9e78e5a9f6179dd91352772ae1a9849246/google_cloud_storage-3.4.1.tar.gz", hash = "sha256:6f041a297e23a4b485fad8c305a7a6e6831855c208bcbe74d00332a909f82268", size = 17238203, upload-time = "2025-10-08T18:43:39.665Z" } wheels = [ @@ -2732,12 +2752,12 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "google-api-core", version = "2.30.3", source = { registry = "https://pypi.org/simple" } }, - { name = "google-auth" }, - { name = "google-cloud-core" }, - { name = "google-crc32c" }, - { name = "google-resumable-media" }, - { name = "requests" }, + { name = "google-api-core", version = "2.30.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14'" }, + { name = "google-auth", marker = "python_full_version < '3.14'" }, + { name = "google-cloud-core", marker = "python_full_version < '3.14'" }, + { name = "google-crc32c", marker = "python_full_version < '3.14'" }, + { name = "google-resumable-media", marker = "python_full_version < '3.14'" }, + { name = "requests", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/4c/47/205eb8e9a1739b5345843e5a425775cbdc472cc38e7eda082ba5b8d02450/google_cloud_storage-3.10.1.tar.gz", hash = "sha256:97db9aa4460727982040edd2bd13ff3d5e2260b5331ad22895802da1fc2a5286", size = 17309950, upload-time = "2026-03-23T09:35:23.409Z" } wheels = [ @@ -3446,7 +3466,8 @@ version = "1.32.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, - { name = "filelock" }, + { name = "filelock", version = "3.20.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "filelock", version = "4.2.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "fsspec" }, { name = "hf-xet", marker = "platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" }, { name = "httpx" }, @@ -4199,13 +4220,13 @@ name = "langchain-classic" version = "1.0.7" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "langchain-core" }, - { name = "langchain-text-splitters" }, - { name = "langsmith" }, - { name = "pydantic" }, - { name = "pyyaml" }, - { name = "requests" }, - { name = "sqlalchemy" }, + { name = "langchain-core", marker = "python_full_version >= '3.11'" }, + { name = "langchain-text-splitters", marker = "python_full_version >= '3.11'" }, + { name = "langsmith", marker = "python_full_version >= '3.11'" }, + { name = "pydantic", marker = "python_full_version >= '3.11'" }, + { name = "pyyaml", marker = "python_full_version >= '3.11'" }, + { name = "requests", marker = "python_full_version >= '3.11'" }, + { name = "sqlalchemy", marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/9b/78/84b5065816f348c39fefa4316f209f0135e8410216340a953bec17d9e4e4/langchain_classic-1.0.7.tar.gz", hash = "sha256:debbec8065e69b95108d2652e8d5c44f4516e19aa8d716c02ed2211c3aee099d", size = 10554118, upload-time = "2026-05-07T15:46:56.8Z" } wheels = [ @@ -4220,18 +4241,18 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "aiohttp" }, - { name = "dataclasses-json" }, - { name = "httpx-sse" }, - { name = "langchain" }, - { name = "langchain-core" }, - { name = "langsmith" }, - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" } }, - { name = "pydantic-settings" }, - { name = "pyyaml" }, - { name = "requests" }, - { name = "sqlalchemy" }, - { name = "tenacity" }, + { name = "aiohttp", marker = "python_full_version < '3.11'" }, + { name = "dataclasses-json", marker = "python_full_version < '3.11'" }, + { name = "httpx-sse", marker = "python_full_version < '3.11'" }, + { name = "langchain", marker = "python_full_version < '3.11'" }, + { name = "langchain-core", marker = "python_full_version < '3.11'" }, + { name = "langsmith", marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pydantic-settings", marker = "python_full_version < '3.11'" }, + { name = "pyyaml", marker = "python_full_version < '3.11'" }, + { name = "requests", marker = "python_full_version < '3.11'" }, + { name = "sqlalchemy", marker = "python_full_version < '3.11'" }, + { name = "tenacity", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/83/49/2ff5354273809e9811392bc24bcffda545a196070666aef27bc6aacf1c21/langchain_community-0.3.31.tar.gz", hash = "sha256:250e4c1041539130f6d6ac6f9386cb018354eafccd917b01a4cff1950b80fd81", size = 33241237, upload-time = "2025-10-07T20:17:57.857Z" } wheels = [ @@ -4249,19 +4270,19 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "aiohttp" }, - { name = "dataclasses-json" }, - { name = "httpx-sse" }, - { name = "langchain-classic" }, - { name = "langchain-core" }, - { name = "langsmith" }, - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "aiohttp", marker = "python_full_version >= '3.11'" }, + { name = "dataclasses-json", marker = "python_full_version >= '3.11'" }, + { name = "httpx-sse", marker = "python_full_version >= '3.11'" }, + { name = "langchain-classic", marker = "python_full_version >= '3.11'" }, + { name = "langchain-core", marker = "python_full_version >= '3.11'" }, + { name = "langsmith", marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "pydantic-settings" }, - { name = "pyyaml" }, - { name = "requests" }, - { name = "sqlalchemy" }, - { name = "tenacity" }, + { name = "pydantic-settings", marker = "python_full_version >= '3.11'" }, + { name = "pyyaml", marker = "python_full_version >= '3.11'" }, + { name = "requests", marker = "python_full_version >= '3.11'" }, + { name = "sqlalchemy", marker = "python_full_version >= '3.11'" }, + { name = "tenacity", marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/53/97/a03585d42b9bdb6fbd935282d6e3348b10322a24e6ce12d0c99eb461d9af/langchain_community-0.4.1.tar.gz", hash = "sha256:f3b211832728ee89f169ddce8579b80a085222ddb4f4ed445a46e977d17b1e85", size = 33241144, upload-time = "2025-10-27T15:20:32.504Z" } wheels = [ @@ -4319,7 +4340,7 @@ name = "langchain-text-splitters" version = "1.1.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "langchain-core" }, + { name = "langchain-core", marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/26/9f/6c545900fefb7b00ddfa3f16b80d61338a0ec68c31c5451eeeab99082760/langchain_text_splitters-1.1.2.tar.gz", hash = "sha256:782a723db0a4746ac91e251c7c1d57fd23636e4f38ed733074e28d7a86f41627", size = 293580, upload-time = "2026-04-16T14:20:39.162Z" } wheels = [ @@ -4581,7 +4602,8 @@ dependencies = [ { name = "certifi" }, { name = "click" }, { name = "fastuuid" }, - { name = "filelock" }, + { name = "filelock", version = "3.20.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "filelock", version = "4.2.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "httpx", extra = ["http2"] }, { name = "huggingface-hub" }, { name = "importlib-metadata" }, @@ -4855,7 +4877,7 @@ requires-dist = [ { name = "fastapi", marker = "extra == 'proxy'", specifier = ">=0.136.3,<1.0" }, { name = "fastapi-sso", marker = "extra == 'proxy'", specifier = ">=0.19.0,<1.0" }, { name = "fastuuid", specifier = ">=0.14.0,<1.0" }, - { name = "filelock", specifier = ">=3.16.1,<4.0" }, + { name = "filelock", specifier = ">=3.16.1,<5" }, { name = "google-cloud-aiplatform", marker = "extra == 'google'", specifier = ">=1.133.0,<2.0" }, { name = "google-cloud-aiplatform", marker = "extra == 'proxy-runtime'", specifier = ">=1.133.0,<2.0" }, { name = "google-cloud-iam", marker = "extra == 'extra-proxy'", specifier = ">=2.19.1,<3.0" }, @@ -5051,10 +5073,10 @@ name = "litellm-admin-mcp" version = "0.1.0" source = { url = "https://github.com/BerriAI/liteadmin-mcp/archive/d35ec9c19c117d4c50cc1eccf6ce1296aac25a1a.tar.gz" } dependencies = [ - { name = "httpx2" }, - { name = "jsonschema" }, - { name = "mcp" }, - { name = "uvicorn" }, + { name = "httpx2", marker = "python_full_version >= '3.12'" }, + { name = "jsonschema", marker = "python_full_version >= '3.12'" }, + { name = "mcp", marker = "python_full_version >= '3.12'" }, + { name = "uvicorn", marker = "python_full_version >= '3.12'" }, ] sdist = { hash = "sha256:86b930f6706fb2da10d53d1798ee7fd14e44fb0afdef0da122cf0e6bb88fa2c3" } @@ -5160,16 +5182,16 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "aiohttp" }, - { name = "chevron" }, - { name = "jsonpickle" }, - { name = "langchain-community", version = "0.3.31", source = { registry = "https://pypi.org/simple" } }, - { name = "packaging" }, - { name = "pydantic" }, - { name = "pyhumps" }, - { name = "requests" }, - { name = "setuptools" }, - { name = "tenacity" }, + { name = "aiohttp", marker = "python_full_version < '3.11'" }, + { name = "chevron", marker = "python_full_version < '3.11'" }, + { name = "jsonpickle", marker = "python_full_version < '3.11'" }, + { name = "langchain-community", version = "0.3.31", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "packaging", marker = "python_full_version < '3.11'" }, + { name = "pydantic", marker = "python_full_version < '3.11'" }, + { name = "pyhumps", marker = "python_full_version < '3.11'" }, + { name = "requests", marker = "python_full_version < '3.11'" }, + { name = "setuptools", marker = "python_full_version < '3.11'" }, + { name = "tenacity", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/4a/6f/9ca1acf766848aaf5f0ac4140c34c91ad0dbfad2654359699644be3352c9/lunary-1.4.36.tar.gz", hash = "sha256:53f002f385c83d9c0e6368e7999923acffbde987f53c5205c2c249c38ee2d75c", size = 20253, upload-time = "2026-02-09T20:49:30.56Z" } wheels = [ @@ -5187,16 +5209,16 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "aiohttp" }, - { name = "chevron" }, - { name = "jsonpickle" }, - { name = "langchain-community", version = "0.4.1", source = { registry = "https://pypi.org/simple" } }, - { name = "packaging" }, - { name = "pydantic" }, - { name = "pyhumps" }, - { name = "requests" }, - { name = "setuptools" }, - { name = "tenacity" }, + { name = "aiohttp", marker = "python_full_version >= '3.11'" }, + { name = "chevron", marker = "python_full_version >= '3.11'" }, + { name = "jsonpickle", marker = "python_full_version >= '3.11'" }, + { name = "langchain-community", version = "0.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "packaging", marker = "python_full_version >= '3.11'" }, + { name = "pydantic", marker = "python_full_version >= '3.11'" }, + { name = "pyhumps", marker = "python_full_version >= '3.11'" }, + { name = "requests", marker = "python_full_version >= '3.11'" }, + { name = "setuptools", marker = "python_full_version >= '3.11'" }, + { name = "tenacity", marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/37/ef/1acbc6957585cc0110e648d787663871717ced3df27fcd3cb5e18fa418f3/lunary-1.4.37.tar.gz", hash = "sha256:1781091e9dceffcc28ebc4be7e085c9fec4102d98d7ca945ed0021e9ce03c36f", size = 20248, upload-time = "2026-02-12T08:15:02.091Z" } wheels = [ @@ -9091,10 +9113,10 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "joblib" }, - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" } }, - { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" } }, - { name = "threadpoolctl" }, + { name = "joblib", marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "threadpoolctl", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/98/c2/a7855e41c9d285dfe86dc50b250978105dce513d6e459ea66a6aeb0e1e0c/scikit_learn-1.7.2.tar.gz", hash = "sha256:20e9e49ecd130598f1ca38a1d85090e1a600147b9c02fa6f15d69cb53d968fda", size = 7193136, upload-time = "2025-09-09T08:21:29.075Z" } wheels = [ @@ -9141,11 +9163,11 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "joblib" }, - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "joblib", marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" } }, - { name = "threadpoolctl" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "threadpoolctl", marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0e/d4/40988bf3b8e34feec1d0e6a051446b1f66225f8529b9309becaeef62b6c4/scikit_learn-1.8.0.tar.gz", hash = "sha256:9bccbb3b40e3de10351f8f5068e105d0f4083b1a65fa07b6634fbc401a6287fd", size = 7335585, upload-time = "2025-12-10T07:08:53.618Z" } wheels = [ @@ -9195,7 +9217,7 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } wheels = [ @@ -9257,7 +9279,7 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } @@ -9342,21 +9364,21 @@ name = "semantic-router" version = "0.1.15" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "aiohttp" }, - { name = "aurelio-sdk" }, - { name = "colorama" }, - { name = "colorlog" }, - { name = "litellm" }, - { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or python_full_version >= '3.14'" }, + { name = "aiohttp", marker = "python_full_version < '3.14'" }, + { name = "aurelio-sdk", marker = "python_full_version < '3.14'" }, + { name = "colorama", marker = "python_full_version < '3.14'" }, + { name = "colorlog", marker = "python_full_version < '3.14'" }, + { name = "litellm", marker = "python_full_version < '3.14'" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and python_full_version < '3.14'" }, - { name = "openai" }, - { name = "pydantic" }, - { name = "pyyaml" }, - { name = "regex" }, - { name = "tiktoken" }, - { name = "tornado", version = "6.5.10", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.14'" }, + { name = "openai", marker = "python_full_version < '3.14'" }, + { name = "pydantic", marker = "python_full_version < '3.14'" }, + { name = "pyyaml", marker = "python_full_version < '3.14'" }, + { name = "regex", marker = "python_full_version < '3.14'" }, + { name = "tiktoken", marker = "python_full_version < '3.14'" }, + { name = "tornado", version = "6.5.10", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "tornado", version = "6.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.14'" }, - { name = "urllib3" }, + { name = "urllib3", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/dc/a9/1a689e916e8b280f1fd8fb335cc059be626a22fe4533baa045d32fcd6de5/semantic_router-0.1.15.tar.gz", hash = "sha256:328256ddc3c2b713101ec69561d6585aecbf1198ea3461e1486289d8c3a35288", size = 95605, upload-time = "2026-05-23T12:58:15.444Z" } wheels = [ @@ -9430,9 +9452,9 @@ name = "smithy-aws-core" version = "0.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "aws-sdk-signers" }, - { name = "smithy-core" }, - { name = "smithy-http" }, + { name = "aws-sdk-signers", marker = "python_full_version >= '3.12'" }, + { name = "smithy-core", marker = "python_full_version >= '3.12'" }, + { name = "smithy-http", marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7d/d3/501c0023548173416109ac42298ca33b708469dc922005770811a597949f/smithy_aws_core-0.11.0.tar.gz", hash = "sha256:29ee89976a520a87e3db557e03e115fdc21a0a60b81161e95174395a1b064da1", size = 38791, upload-time = "2026-08-24T21:16:59.631Z" } wheels = [ @@ -9441,10 +9463,10 @@ wheels = [ [package.optional-dependencies] eventstream = [ - { name = "smithy-aws-event-stream" }, + { name = "smithy-aws-event-stream", marker = "python_full_version >= '3.12'" }, ] json = [ - { name = "smithy-json" }, + { name = "smithy-json", marker = "python_full_version >= '3.12'" }, ] [[package]] @@ -9452,7 +9474,7 @@ name = "smithy-aws-event-stream" version = "0.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "smithy-core" }, + { name = "smithy-core", marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/38/0e/6efb3a4ed92c0f1ada6de060ac92e7115a1e34d0ab1fb99a6056734a88ea/smithy_aws_event_stream-0.3.0.tar.gz", hash = "sha256:a0e227367a973144e205a075d0a424f95c92f26656a1018d08900da2ae547c49", size = 12818, upload-time = "2026-05-05T18:04:14.317Z" } wheels = [ @@ -9473,7 +9495,7 @@ name = "smithy-http" version = "0.5.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "smithy-core" }, + { name = "smithy-core", marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/98/78/b5f3113d6c8f0bc1f9777a7f5ca84b892d29efac05850e14f7d4f7e645b5/smithy_http-0.5.0.tar.gz", hash = "sha256:bb4a19672f7c7eeb872a308f777eb505281a5bafb1ee3d1ea9c760c06c352510", size = 31122, upload-time = "2026-08-24T21:16:56.488Z" } wheels = [ @@ -9482,11 +9504,11 @@ wheels = [ [package.optional-dependencies] aiohttp = [ - { name = "aiohttp" }, - { name = "yarl" }, + { name = "aiohttp", marker = "python_full_version >= '3.12'" }, + { name = "yarl", marker = "python_full_version >= '3.12'" }, ] awscrt = [ - { name = "awscrt" }, + { name = "awscrt", marker = "python_full_version >= '3.12'" }, ] [[package]] @@ -9494,8 +9516,8 @@ name = "smithy-json" version = "0.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ijson" }, - { name = "smithy-core" }, + { name = "ijson", marker = "python_full_version >= '3.12'" }, + { name = "smithy-core", marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c7/ac/04164eefb3da7479f52f6535b4b39cc8384c292cb2bb74279f2acc4f4b4d/smithy_json-0.3.0.tar.gz", hash = "sha256:c81c7034587e01bc64767cbbecb05a7d65ca9070612fd94e8a03e80540290a22", size = 7956, upload-time = "2026-08-20T17:55:32.177Z" } wheels = [ @@ -9573,23 +9595,23 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "alabaster" }, - { name = "babel" }, - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "docutils", version = "0.21.2", source = { registry = "https://pypi.org/simple" } }, - { name = "imagesize" }, - { name = "jinja2" }, - { name = "packaging" }, - { name = "pygments" }, - { name = "requests" }, - { name = "snowballstemmer" }, - { name = "sphinxcontrib-applehelp" }, - { name = "sphinxcontrib-devhelp" }, - { name = "sphinxcontrib-htmlhelp" }, - { name = "sphinxcontrib-jsmath" }, - { name = "sphinxcontrib-qthelp" }, - { name = "sphinxcontrib-serializinghtml" }, - { name = "tomli" }, + { name = "alabaster", marker = "python_full_version < '3.11'" }, + { name = "babel", marker = "python_full_version < '3.11'" }, + { name = "colorama", marker = "python_full_version < '3.11' and sys_platform == 'win32'" }, + { name = "docutils", version = "0.21.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "imagesize", marker = "python_full_version < '3.11'" }, + { name = "jinja2", marker = "python_full_version < '3.11'" }, + { name = "packaging", marker = "python_full_version < '3.11'" }, + { name = "pygments", marker = "python_full_version < '3.11'" }, + { name = "requests", marker = "python_full_version < '3.11'" }, + { name = "snowballstemmer", marker = "python_full_version < '3.11'" }, + { name = "sphinxcontrib-applehelp", marker = "python_full_version < '3.11'" }, + { name = "sphinxcontrib-devhelp", marker = "python_full_version < '3.11'" }, + { name = "sphinxcontrib-htmlhelp", marker = "python_full_version < '3.11'" }, + { name = "sphinxcontrib-jsmath", marker = "python_full_version < '3.11'" }, + { name = "sphinxcontrib-qthelp", marker = "python_full_version < '3.11'" }, + { name = "sphinxcontrib-serializinghtml", marker = "python_full_version < '3.11'" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6f/6d/be0b61178fe2cdcb67e2a92fc9ebb488e3c51c4f74a36a7824c0adf23425/sphinx-8.1.3.tar.gz", hash = "sha256:43c1911eecb0d3e161ad78611bc905d1ad0e523e4ddc202a58a821773dc4c927", size = 8184611, upload-time = "2024-10-13T20:27:13.93Z" } wheels = [ @@ -9604,23 +9626,23 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "alabaster" }, - { name = "babel" }, - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "docutils", version = "0.22.4", source = { registry = "https://pypi.org/simple" } }, - { name = "imagesize" }, - { name = "jinja2" }, - { name = "packaging" }, - { name = "pygments" }, - { name = "requests" }, - { name = "roman-numerals" }, - { name = "snowballstemmer" }, - { name = "sphinxcontrib-applehelp" }, - { name = "sphinxcontrib-devhelp" }, - { name = "sphinxcontrib-htmlhelp" }, - { name = "sphinxcontrib-jsmath" }, - { name = "sphinxcontrib-qthelp" }, - { name = "sphinxcontrib-serializinghtml" }, + { name = "alabaster", marker = "python_full_version == '3.11.*'" }, + { name = "babel", marker = "python_full_version == '3.11.*'" }, + { name = "colorama", marker = "python_full_version == '3.11.*' and sys_platform == 'win32'" }, + { name = "docutils", version = "0.22.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "imagesize", marker = "python_full_version == '3.11.*'" }, + { name = "jinja2", marker = "python_full_version == '3.11.*'" }, + { name = "packaging", marker = "python_full_version == '3.11.*'" }, + { name = "pygments", marker = "python_full_version == '3.11.*'" }, + { name = "requests", marker = "python_full_version == '3.11.*'" }, + { name = "roman-numerals", marker = "python_full_version == '3.11.*'" }, + { name = "snowballstemmer", marker = "python_full_version == '3.11.*'" }, + { name = "sphinxcontrib-applehelp", marker = "python_full_version == '3.11.*'" }, + { name = "sphinxcontrib-devhelp", marker = "python_full_version == '3.11.*'" }, + { name = "sphinxcontrib-htmlhelp", marker = "python_full_version == '3.11.*'" }, + { name = "sphinxcontrib-jsmath", marker = "python_full_version == '3.11.*'" }, + { name = "sphinxcontrib-qthelp", marker = "python_full_version == '3.11.*'" }, + { name = "sphinxcontrib-serializinghtml", marker = "python_full_version == '3.11.*'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/50/a8c6ccc36d5eacdfd7913ddccd15a9cee03ecafc5ee2bc40e1f168d85022/sphinx-9.0.4.tar.gz", hash = "sha256:594ef59d042972abbc581d8baa577404abe4e6c3b04ef61bd7fc2acbd51f3fa3", size = 8710502, upload-time = "2025-12-04T07:45:27.343Z" } wheels = [ @@ -9637,23 +9659,23 @@ resolution-markers = [ "python_full_version == '3.12.*'", ] dependencies = [ - { name = "alabaster" }, - { name = "babel" }, - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "docutils", version = "0.22.4", source = { registry = "https://pypi.org/simple" } }, - { name = "imagesize" }, - { name = "jinja2" }, - { name = "packaging" }, - { name = "pygments" }, - { name = "requests" }, - { name = "roman-numerals" }, - { name = "snowballstemmer" }, - { name = "sphinxcontrib-applehelp" }, - { name = "sphinxcontrib-devhelp" }, - { name = "sphinxcontrib-htmlhelp" }, - { name = "sphinxcontrib-jsmath" }, - { name = "sphinxcontrib-qthelp" }, - { name = "sphinxcontrib-serializinghtml" }, + { name = "alabaster", marker = "python_full_version >= '3.12'" }, + { name = "babel", marker = "python_full_version >= '3.12'" }, + { name = "colorama", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" }, + { name = "docutils", version = "0.22.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "imagesize", marker = "python_full_version >= '3.12'" }, + { name = "jinja2", marker = "python_full_version >= '3.12'" }, + { name = "packaging", marker = "python_full_version >= '3.12'" }, + { name = "pygments", marker = "python_full_version >= '3.12'" }, + { name = "requests", marker = "python_full_version >= '3.12'" }, + { name = "roman-numerals", marker = "python_full_version >= '3.12'" }, + { name = "snowballstemmer", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-applehelp", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-devhelp", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-htmlhelp", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-jsmath", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-qthelp", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-serializinghtml", marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cd/bd/f08eb0f4eed5c83f1ba2a3bd18f7745a2b1525fad70660a1c00224ec468a/sphinx-9.1.0.tar.gz", hash = "sha256:7741722357dd75f8190766926071fed3bdc211c74dd2d7d4df5404da95930ddb", size = 8718324, upload-time = "2025-12-31T15:09:27.646Z" } wheels = [ @@ -9801,8 +9823,8 @@ name = "standard-aifc" version = "3.13.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "audioop-lts" }, - { name = "standard-chunk" }, + { name = "audioop-lts", marker = "python_full_version >= '3.13'" }, + { name = "standard-chunk", marker = "python_full_version >= '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c4/53/6050dc3dde1671eb3db592c13b55a8005e5040131f7509cef0215212cb84/standard_aifc-3.13.0.tar.gz", hash = "sha256:64e249c7cb4b3daf2fdba4e95721f811bde8bdfc43ad9f936589b7bb2fae2e43", size = 15240, upload-time = "2024-10-30T16:01:31.772Z" } wheels = [ @@ -9823,7 +9845,7 @@ name = "standard-sunau" version = "3.13.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "audioop-lts" }, + { name = "audioop-lts", marker = "python_full_version >= '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/66/e3/ce8d38cb2d70e05ffeddc28bb09bad77cfef979eb0a299c9117f7ed4e6a9/standard_sunau-3.13.0.tar.gz", hash = "sha256:b319a1ac95a09a2378a8442f403c66f4fd4b36616d6df6ae82b8e536ee790908", size = 9368, upload-time = "2024-10-30T16:01:41.626Z" } wheels = [ @@ -9857,8 +9879,8 @@ name = "taskgroup" version = "0.2.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "exceptiongroup" }, - { name = "typing-extensions" }, + { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f0/8d/e218e0160cc1b692e6e0e5ba34e8865dbb171efeb5fc9a704544b3020605/taskgroup-0.2.2.tar.gz", hash = "sha256:078483ac3e78f2e3f973e2edbf6941374fbea81b9c5d0a96f51d297717f4752d", size = 11504, upload-time = "2025-01-03T09:24:13.761Z" } wheels = [ From 3867d48dd56851dde781914c5786021eb1154ad0 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Tue, 20 Oct 2026 16:00:00 -0700 Subject: [PATCH 10/41] fix(mcp): backport #47034 to lts/1.108.x Upgrade to MCP SDK 2.3.0 and raise the per-event SSE size cap so MCP tool results over 1 MiB stop failing with "SSE stream ended without a response" (cherry picked from commit 4aec80f5ee1b3bbd2804abc123fe6663be77f775) Co-Authored-By: Claude Opus 5.5 --- litellm/constants.py | 3 + litellm/experimental_mcp_client/client.py | 2 + pyproject.toml | 11 +-- tests/integration/_support/mcp.py | 6 +- tests/integration/mcp/test_mcp_transports.py | 72 +++++++++++++++++++ .../test_mcp_client.py | 4 +- tests/unit/test_constants.py | 28 ++++++++ uv.lock | 25 +++---- 8 files changed, 131 insertions(+), 20 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 9a994cd30b4..bfdb605d232 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -237,6 +237,9 @@ MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS: Final = get_env_int("MCP_PER_USER_TOKE # MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers. MCP_CLIENT_TIMEOUT: Final = get_env_float("LITELLM_MCP_CLIENT_TIMEOUT", 60.0) +MCP_MAX_SSE_EVENT_SIZE_BYTES: Final = get_env_int_in_range( + "LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES", 64 * 1024 * 1024, 1, sys.maxsize +) MCP_TOOL_LISTING_TIMEOUT: Final = get_env_float("LITELLM_MCP_TOOL_LISTING_TIMEOUT", 30.0) MCP_METADATA_TIMEOUT: Final = get_env_float("LITELLM_MCP_METADATA_TIMEOUT", 10.0) MCP_HEALTH_CHECK_TIMEOUT: Final = get_env_float("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", 10.0) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index e2057e4b627..c0d7824682a 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -95,6 +95,7 @@ from pydantic import AnyUrl, TypeAdapter from litellm._logging import verbose_logger from litellm.constants import ( MCP_CLIENT_TIMEOUT, + MCP_MAX_SSE_EVENT_SIZE_BYTES, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_MAX_PAGES, MCP_TOOL_LISTING_TIMEOUT, @@ -631,6 +632,7 @@ class MCPClient: transport_ctx: Final = streamable_http_client( url=self.server_url, http_client=http_client, + max_sse_event_size=MCP_MAX_SSE_EVENT_SIZE_BYTES, ) return transport_ctx, http_client diff --git a/pyproject.toml b/pyproject.toml index fca9ea6c881..f166279574d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,8 +75,8 @@ proxy = [ "boto3>=1.43.117,<2.0", "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", - "mcp>=2.2.0,<3", - "httpx2>=2.5.0,<3", + "mcp>=2.3.0,<3", + "httpx2>=2.10.0,<3", "pydantic>=2.12.0,<3", "litellm-proxy-extras==0.4.110", "litellm-enterprise==0.1.79.post1", @@ -124,7 +124,7 @@ utils = [ "numpydoc>=1.8.0,<2.0", ] caching = ["diskcache>=5.6.3,<6.0"] -mcp = ["mcp>=2.2.0,<3", "httpx2>=2.5.0,<3", "pydantic>=2.12.0,<3"] +mcp = ["mcp>=2.3.0,<3", "httpx2>=2.10.0,<3", "pydantic>=2.12.0,<3"] # Driver for the MongoDB Atlas vector store; Atlas Vector Search has no HTTP query API. # The floor is 4.9 because that is the release AsyncMongoClient landed in. # SAML SSO for the admin UI. python3-saml pulls in xmlsec/lxml, whose wheels @@ -242,7 +242,7 @@ e2e-dev = [ "locust==2.45.0", "anthropic==0.84.0", "psutil==7.2.2", - "mcp>=2.2.0,<3", + "mcp>=2.3.0,<3", ] admin-mcp = [ "litellm-admin-mcp @ https://github.com/BerriAI/liteadmin-mcp/archive/d35ec9c19c117d4c50cc1eccf6ce1296aac25a1a.tar.gz#sha256=86b930f6706fb2da10d53d1798ee7fd14e44fb0afdef0da122cf0e6bb88fa2c3 ; python_version >= '3.12'", @@ -302,7 +302,7 @@ healthcheck = [ benchmarks = [ "pytest==9.0.3", "pytest-codspeed==4.3.0", - "mcp>=2.2.0,<3", + "mcp>=2.3.0,<3", "a2a-sdk==1.1.0", ] @@ -363,6 +363,7 @@ override-dependencies = [ # a2a-sdk 1.x requires packaging>=24.0; lunary 1.4.x still caps at <24.0. "packaging>=24.0", "cryptography>=50.0.0,<51.0", + "mcp>=2.3.0,<3", ] default-groups = ["dev"] required-version = ">=0.10.9" diff --git a/tests/integration/_support/mcp.py b/tests/integration/_support/mcp.py index f69079ce90b..99bf71cd459 100644 --- a/tests/integration/_support/mcp.py +++ b/tests/integration/_support/mcp.py @@ -282,6 +282,7 @@ class ScriptedTool: respond: Callable[[JsonRpc], Reply | JsonRpc] description: str | Callable[[Mapping[str, str]], str] | None = None input_schema: JsonRpc = field(default_factory=lambda: {"type": "object"}) + event_stream: bool = False def listing(self, headers: Mapping[str, str]) -> JsonRpc: described: Final = self.description(headers) if callable(self.description) else self.description @@ -334,7 +335,10 @@ def scripted_peer(*tools: ScriptedTool) -> Iterator[McpPeer]: if tool is None: return jsonrpc_error(identity, -32602, "unknown tool") produced: Final = tool.respond(body["params"]) - return produced if isinstance(produced, Reply) else jsonrpc_reply(identity, produced) + reply: Final = produced if isinstance(produced, Reply) else jsonrpc_reply(identity, produced) + if not tool.event_stream: + return reply + return Reply(content_type="text/event-stream", body=b"event: message\ndata: " + reply.body + b"\n\n") with wire_server(provider) as wire: yield McpPeer(wire.url + "/mcp", observed) diff --git a/tests/integration/mcp/test_mcp_transports.py b/tests/integration/mcp/test_mcp_transports.py index 928aa9ff121..ef240cba570 100644 --- a/tests/integration/mcp/test_mcp_transports.py +++ b/tests/integration/mcp/test_mcp_transports.py @@ -2,10 +2,12 @@ import builtins import json import uuid from collections.abc import Mapping +from pathlib import Path from typing import Final import httpx import pytest +import yaml from mcp.shared.exceptions import MCPError from integration._support.client import Gateway from integration._support.database import read_rows @@ -17,15 +19,21 @@ from integration._support.mcp import ( McpCaller, Outcome, PeerKind, + ScriptedTool, _outcome_from_rpc, _parse_rpc_body, + call_tool, + listed_tools, mcp_peer, official_client_outcomes, peer_of, register_mcp, + scripted_peer, + text_result, tool_calls, ) from integration._support.mcp_grants import create_toolset +from integration._support.process import owned_proxy ADD: Final = {"http": "add", "sse": "add", "stdio": "add", "openapi": "getpet"} ARGUMENTS: Final = {"add": {"a": 3, "b": 4}, "getpet": {"petId": "7"}} @@ -817,3 +825,67 @@ def test_sdk_progress_and_cancellation_reach_upstream_without_losing_session( ), ) as gateway: asyncio.run(exercise(gateway)) + + +def test_tool_result_larger_than_one_mib_in_a_single_sse_event_reaches_the_caller(gateway: Gateway) -> None: + payload: Final = "x" * (3 * 1024 * 1024) + tool: Final = ScriptedTool("dump", lambda _: text_result(payload), event_stream=True) + + with scripted_peer(tool) as peer, gateway.scenario() as scenario: + alias: Final = "bigsse" + uuid.uuid4().hex[:8] + identity: Final = register_mcp(scenario, peer, alias) + key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) + name: Final = next(full for full in listed_tools(gateway, key, identity) if full.endswith("dump")) + response: Final = call_tool(gateway, key, identity, name, {}) + assert response.status_code == 200, response.text[:500] + body: Final = response.json() + assert body["isError"] is False, response.text[:500] + assert body["content"][0]["text"] == payload + + +def test_configured_sse_limit_rejects_large_results_and_allows_the_next_call(tmp_path: Path) -> None: + limit: Final = 4096 + payload: Final = "x" * 512 + small: Final = ScriptedTool("small", lambda _: text_result(payload), event_stream=True) + large: Final = ScriptedTool("large", lambda _: text_result("x" * (2 * limit)), event_stream=True) + key: Final = "sk-local-mcp-sse-limit-regression-key" + with scripted_peer(small, large) as peer, httpx.Client() as client: + config: Final = tmp_path / "sse-limit.yaml" + config.write_text( + yaml.safe_dump( + { + "general_settings": {"master_key": key, "store_model_in_db": False}, + "mcp_servers": {"limited": peer.registration()}, + } + ) + ) + seed: Final = Gateway(client, key, peer.url) + with owned_proxy( + seed, + tmp_path / "proxy", + { + "STORE_MODEL_IN_DB": "False", + "DISABLE_SCHEMA_UPDATE": "true", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES": str(limit), + }, + config=config, + database_setup=(), + remove_environment=( + "DATABASE_URL", + "DATABASE_URL_READ_REPLICA", + "DIRECT_URL", + "INTEGRATION_PROXY_DATABASE_URL", + "INTEGRATION_PROXY_READ_REPLICA_URL", + ), + ) as gateway: + assert set(listed_tools(gateway, key)) == {"small", "large"} + for name in ("small", "large", "small"): + response: Final = call_tool(gateway, key, "limited", name, {}) + assert response.status_code == 200, response.text[:500] + body: Final = response.json() + assert body["isError"] is (name == "large"), response.text[:500] + if name == "large": + assert str(limit) in body["content"][0]["text"], response.text[:500] + else: + assert body["content"][0]["text"] == payload, response.text[:500] diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index 418738a9dca..2787c25a6f9 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -1503,8 +1503,8 @@ def test_mcp_extra_matches_proxy_extra_and_supports_streamable_http(): } specifier: Final = Requirement(mcp_extra["mcp"]).specifier - assert not specifier.contains("1.28.1") - assert specifier.contains("2.2.0") + assert not specifier.contains("2.2.0") + assert specifier.contains("2.3.0") with (pyproject_path.parent / "uv.lock").open("rb") as f: locked = tomllib.load(f) mcp_versions: Final = [package["version"] for package in locked["package"] if package["name"] == "mcp"] diff --git a/tests/unit/test_constants.py b/tests/unit/test_constants.py index 0d0e77765b7..99b049f1803 100644 --- a/tests/unit/test_constants.py +++ b/tests/unit/test_constants.py @@ -1,6 +1,7 @@ import ast, asyncio, os import inspect import json +import sys from unittest import mock import httpx @@ -117,6 +118,33 @@ def test_cli_jwt_expiration_hours_from_environment( importlib.reload(litellm.constants) +@pytest.mark.parametrize( + ("raw", "expected"), + ( + (None, 64 * 1024 * 1024), + ("", 64 * 1024 * 1024), + ("invalid", 64 * 1024 * 1024), + ("0", 64 * 1024 * 1024), + ("-1", 64 * 1024 * 1024), + ("1", 1), + ("200000", 200000), + (str(sys.maxsize), sys.maxsize), + (str(sys.maxsize + 1), 64 * 1024 * 1024), + ), +) +def test_mcp_sse_event_limit_uses_positive_override_or_default(raw: str | None, expected: int) -> None: + try: + with pytest.MonkeyPatch.context() as environment: + if raw is None: + environment.delenv("LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES", raising=False) + else: + environment.setenv("LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES", raw) + importlib.reload(constants) + assert constants.MCP_MAX_SSE_EVENT_SIZE_BYTES == expected + finally: + importlib.reload(constants) + + def _is_blank(value: object) -> bool: return isinstance(value, str) and not value.strip() diff --git a/uv.lock b/uv.lock index 0560d07f501..cee34c63b81 100644 --- a/uv.lock +++ b/uv.lock @@ -42,6 +42,7 @@ constraints = [ ] overrides = [ { name = "cryptography", specifier = ">=50.0.0,<51.0" }, + { name = "mcp", specifier = ">=2.3.0,<3" }, { name = "packaging", specifier = ">=24.0" }, ] @@ -4869,8 +4870,8 @@ requires-dist = [ { name = "gunicorn", marker = "extra == 'proxy'", specifier = ">=23.0.0,<24.0" }, { name = "hiredis", marker = "extra == 'proxy'", specifier = ">=3.0.0,<4.0" }, { name = "httpx", extras = ["http2"], specifier = ">=0.28.0,<1.0" }, - { name = "httpx2", marker = "extra == 'mcp'", specifier = ">=2.5.0,<3" }, - { name = "httpx2", marker = "extra == 'proxy'", specifier = ">=2.5.0,<3" }, + { name = "httpx2", marker = "extra == 'mcp'", specifier = ">=2.10.0,<3" }, + { name = "httpx2", marker = "extra == 'proxy'", specifier = ">=2.10.0,<3" }, { name = "huggingface-hub", specifier = ">=0.34.0,<2.0" }, { name = "importlib-metadata", specifier = ">=8.0.0,<9.0" }, { name = "inquirerpy", marker = "extra == 'cli'", specifier = ">=0.3.4,<1.0" }, @@ -4883,8 +4884,8 @@ requires-dist = [ { name = "litellm-proxy-extras", marker = "extra == 'proxy'", editable = "litellm-proxy-extras" }, { name = "llm-sandbox", marker = "extra == 'proxy-runtime'", specifier = ">=0.3.39,<1.0" }, { name = "mangum", marker = "extra == 'proxy-runtime'", specifier = ">=0.17.0,<1.0" }, - { name = "mcp", marker = "extra == 'mcp'", specifier = ">=2.2.0,<3" }, - { name = "mcp", marker = "extra == 'proxy'", specifier = ">=2.2.0,<3" }, + { name = "mcp", marker = "extra == 'mcp'", specifier = ">=2.3.0,<3" }, + { name = "mcp", marker = "extra == 'proxy'", specifier = ">=2.3.0,<3" }, { name = "mlflow", marker = "extra == 'mlflow'", specifier = ">=3.17.2,<4.0" }, { name = "numpy", marker = "extra == 'stt-nvidia-riva'", specifier = ">=1.26.0" }, { name = "numpydoc", marker = "extra == 'utils'", specifier = ">=1.8.0,<2.0" }, @@ -4944,7 +4945,7 @@ provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "mcp", "sa admin-mcp = [{ name = "litellm-admin-mcp", marker = "python_full_version >= '3.12'", url = "https://github.com/BerriAI/liteadmin-mcp/archive/d35ec9c19c117d4c50cc1eccf6ce1296aac25a1a.tar.gz" }] benchmarks = [ { name = "a2a-sdk", specifier = "==1.1.0" }, - { name = "mcp", specifier = ">=2.2.0,<3" }, + { name = "mcp", specifier = ">=2.3.0,<3" }, { name = "pytest", specifier = "==9.0.3" }, { name = "pytest-codspeed", specifier = "==4.3.0" }, ] @@ -5023,7 +5024,7 @@ dev = [ e2e-dev = [ { name = "anthropic", specifier = "==0.84.0" }, { name = "locust", specifier = "==2.45.0" }, - { name = "mcp", specifier = ">=2.2.0,<3" }, + { name = "mcp", specifier = ">=2.3.0,<3" }, { name = "playwright", specifier = "==1.61.0" }, { name = "psutil", specifier = "==7.2.2" }, { name = "websockets", specifier = ">=15.0.1,<16.0" }, @@ -5541,7 +5542,7 @@ wheels = [ [[package]] name = "mcp" -version = "2.2.0" +version = "2.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -5559,22 +5560,22 @@ dependencies = [ { name = "typing-inspection" }, { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/76/31/ac54fb0fdd5b37de704486e288bba4fbbb463f24cfcfedbede407b854513/mcp-2.2.0.tar.gz", hash = "sha256:2dc37ecb1974becdcebdbf7561e7c15a07dbbf20ba21ba16c3593b3038b3afbd", size = 4084129, upload-time = "2026-09-07T16:06:23.439Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9d/8d/e0d339616f4810e9051d4aba6887afab289ab1f81875fe908b606cdfd0e3/mcp-2.3.0.tar.gz", hash = "sha256:8b147a50441cf059dc88c684e0aeed3687f0aa0f39c6cde7b90330effd2b34d8", size = 4171955, upload-time = "2026-10-02T22:06:56.091Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1b/ff/8e7eade68b8a28f7da0ed1085544341b51f9c935dbf6b95c76b7edfea6a0/mcp-2.2.0-py3-none-any.whl", hash = "sha256:bde982589473a060ae145e3406e9a5333fe538c97229ba841f5a7f92be004f81", size = 365656, upload-time = "2026-09-07T16:06:19.711Z" }, + { url = "https://files.pythonhosted.org/packages/20/d7/4a2878f128df729db3ca04e71bfce26405cc2e0c9ebe19c592feeb4c61b9/mcp-2.3.0-py3-none-any.whl", hash = "sha256:dd0c44c089d16453e8ae31a3877a0054d7a2314caaa81f5e0541b9b1734b2377", size = 368664, upload-time = "2026-10-02T22:06:52.119Z" }, ] [[package]] name = "mcp-types" -version = "2.2.0" +version = "2.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pydantic" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ae/91/762d7755d971aff8a28d75f7961656148edf27875c8026e6385aaab08ae7/mcp_types-2.2.0.tar.gz", hash = "sha256:d3ed53703ddd10d9c6399f29d322bb66f3f67ab41348ac8556ba23e07fedefad", size = 65892, upload-time = "2026-09-07T16:06:25.187Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9e/2d/7c251e34207f6c51000312fc8839111ac45cfe02023f90b44e7f1051dd8e/mcp_types-2.3.0.tar.gz", hash = "sha256:d1e46549edb35ee19a94940fcee6d1addd7e589ab7ea92dda83f5d84781fc362", size = 65898, upload-time = "2026-10-02T22:06:57.863Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/8f/d7/6ffba5d8cd5dd9b8a19478875c50e04945314ba5074e84d749283f27f62d/mcp_types-2.2.0-py3-none-any.whl", hash = "sha256:ea476b73ee86709ab5abc9452385ed36cc05907e582355622e294595c9a04f13", size = 69106, upload-time = "2026-09-07T16:06:21.461Z" }, + { url = "https://files.pythonhosted.org/packages/ce/48/32583066b155bb5d4507e7887dc88ee23118cc4de8e2a12870c49719e2a8/mcp_types-2.3.0-py3-none-any.whl", hash = "sha256:968efdbdaedfab06adae40d378a34395f1090c5921d4be3c9cde283aaf76d91d", size = 69108, upload-time = "2026-10-02T22:06:54.099Z" }, ] [[package]] From 598138fdf052bf8fcf39df6d6391ddf4fcd59995 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 7 Nov 2026 10:23:22 +0900 Subject: [PATCH 11/41] chore(deps): widen litellm-core filelock cap to <5 to match the root manifest --- packaging/litellm-core/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packaging/litellm-core/pyproject.toml b/packaging/litellm-core/pyproject.toml index 0914deb074d..84d428368b7 100644 --- a/packaging/litellm-core/pyproject.toml +++ b/packaging/litellm-core/pyproject.toml @@ -11,7 +11,7 @@ authors = [ ] dependencies = [ "fastuuid>=0.14.0,<1.0", - "filelock>=3.16.1,<4.0", + "filelock>=3.16.1,<5", "httpx[http2]>=0.28.0,<1.0", "openai>=2.20.0,<4.0.0", "python-dateutil>=2.8.2,<3.0", From 048b587b0c590d5ec3aeace987368df83d2390b8 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 4 Nov 2026 09:55:53 -0800 Subject: [PATCH 12/41] test(e2e): retire the quiet_stack lock and batch the claude_code model registration (#49984) * test(e2e): retire the quiet_stack lock quiet_stack held an exclusive host-wide lock for six tests, draining every other xdist worker while they ran. On the full e2e lane that left 35-38% of worker time idle. Each of those tests now lives where it can run without pausing the suite - the two memory tests move to perf/ behind the perf marker, so the perf lane runs them on its own fresh stack. Their refusing deployments use a literal key because the perf stack carries no provider keys - the stored-request size check from the memory test becomes an integration test on an owned proxy that stores prompts - the discount and margin spend-row tests are dropped from e2e. Their breakdown assertions join the existing integration pricing test, which already owns the discount and margin repricing flow - the Bedrock guardrail cost test stays in e2e without the pricing reset fixture, since nothing else on the stack writes pricing config now - the MCP semantic tool filter test turns on a proxy-wide setting, so it gets a new proxy_wide_config opt-in marker and runs after the main run * test(e2e): register the claude_code deployments as one batch The claude_code session fixture registered its 24 deployments one at a time on every xdist worker, and each one waited out the 15 s propagation window on its own, about 6 minutes per worker. ProxyClient.register_models posts the whole batch, polls each model until every replica serves it, then waits out propagation once from the last write. A model that fails to post or never becomes servable comes back as ModelNotRegistered, and an unservable one is deleted, so the fixture keeps reporting per-deployment failures the way it did (cherry picked from commit b18188dee2044de608d6b4389bbff7cf13366ff4) --- tests/e2e/AGENTS.md | 8 +- tests/e2e/claude_code/conftest.py | 35 ++-- tests/e2e/conftest.py | 21 +-- .../coverage_registry/quota_management.yaml | 2 - tests/e2e/e2e_config.py | 2 +- .../mcp/test_mcp_semantic_tool_filter_e2e.py | 5 +- tests/e2e/models.py | 21 --- tests/e2e/perf/test_perf_cases.py | 2 +- .../test_reliability_memory_e2e.py | 111 ++++++------- tests/e2e/proxy_client.py | 96 +++++++---- tests/e2e/pytest.ini | 2 +- ...2e.py => test_guardrail_cost_spend_e2e.py} | 152 +----------------- tests/e2e/stack_lock.py | 45 ------ tests/e2e_harness/AGENTS.md | 2 +- tests/e2e_harness/test_proxy_client.py | 90 ++++++++++- tests/e2e_harness/test_stack_lock.py | 117 -------------- .../pricing/test_cost_margin_discount.py | 32 +++- .../test_failing_request_stored_snapshot.py | 84 ++++++++++ 18 files changed, 358 insertions(+), 469 deletions(-) rename tests/e2e/{router => perf}/test_reliability_memory_e2e.py (77%) rename tests/e2e/quota_management/spend_tracking/{test_pricing_config_spend_e2e.py => test_guardrail_cost_spend_e2e.py} (52%) delete mode 100644 tests/e2e/stack_lock.py delete mode 100644 tests/e2e_harness/test_stack_lock.py create mode 100644 tests/integration/spend/test_failing_request_stored_snapshot.py diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index d38dafedf75..53e203a4d7b 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -40,11 +40,11 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) - `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface, not Playwright) - `a2a/` - the A2A (agent-to-agent) surface: admin registration via `/v1/agents`, proxy-fronted card discovery at `/.well-known/agent-card.json`, and JSON-RPC `message/send` invocation, driving agents backed by the litellm completion bridge (a real provider) and asserting protocol-version normalization (0.3 vs 1.0) -- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below), plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions` in `test_mcp_chat_completion_oauth_e2e.py` and direct MCP protocol operations in `test_mcp_oauth_happy_path_e2e.py`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion or protocol call lists and executes the server's tools with the stored per-user token +- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below), plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions` in `test_mcp_chat_completion_oauth_e2e.py` and direct MCP protocol operations in `test_mcp_oauth_happy_path_e2e.py`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion or protocol call lists and executes the server's tools with the stored per-user token. The semantic tool filter test turns on a proxy-wide setting, so it is marked `proxy_wide_config` and runs on the redeployed managed-files stack after the main run, never next to it - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection -- `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) -- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares) and `Limit`s (one metric at one percentile under one ceiling). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` +- `router/` - routing and reliability behavior (fallbacks, cooldowns) +- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). The memory tests (`test_reliability_memory_e2e.py`, run by the perf lane on its own fresh stack: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches, the check for the v1.100.0 retry-breadcrumb leak). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares) and `Limit`s (one metric at one percentile under one ceiling). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests @@ -98,7 +98,7 @@ Each suite provides its own `client` fixture (see `llm_translation/passthrough_c Request and response bodies are typed pydantic models in `models.py`. Only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass -Mark live tests with `@pytest.mark.e2e` (on the class or the module). Coverage of the harness itself lives outside the suite in `tests/e2e_harness/` (see its `AGENTS.md`) and runs without a proxy, so nothing under `tests/e2e/` is a markerless test. Add `@pytest.mark.quiet_stack` to a test that measures the proxy itself (RSS, latency): the shared stack lock in `stack_lock.py` then runs it while no other test on the host is hitting the stack, marked or not, so the reading depends only on the test's own traffic. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache +Mark live tests with `@pytest.mark.e2e` (on the class or the module). Coverage of the harness itself lives outside the suite in `tests/e2e_harness/` (see its `AGENTS.md`) and runs without a proxy, so nothing under `tests/e2e/` is a markerless test. A test that measures the proxy itself (RSS, latency) cannot share the stack with the rest of the suite, so it goes in `perf/`, marked `perf`. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache ## Record and replay fixtures diff --git a/tests/e2e/claude_code/conftest.py b/tests/e2e/claude_code/conftest.py index b2442d492cf..14201ef6846 100644 --- a/tests/e2e/claude_code/conftest.py +++ b/tests/e2e/claude_code/conftest.py @@ -597,14 +597,21 @@ def _build_control_plane_client(proxy_config: ProxyConfig): ) -def _register_deployment(proxy, deployment: CompatDeployment) -> str: - """Register one deployment and return its proxy-assigned model_id - once it is servable on the data plane. The aliases are shared by every - cell and, under xdist, by every worker, so no call to them belongs to - one test and none is cached: the matrix exists to reach real providers.""" - return proxy.create_model( - deployment.model_name, - deployment.litellm_params, +def _register_deployments(proxy, deployments: tuple[CompatDeployment, ...]): + """The aliases are shared by every cell and, under xdist, by every worker, + so no call to them belongs to one test and none is cached: the matrix + exists to reach real providers.""" + from models import ModelInfoBody, ModelNewBody + + return proxy.register_models( + tuple( + ModelNewBody( + model_name=deployment.model_name, + litellm_params=deployment.litellm_params, + model_info=ModelInfoBody(), + ) + for deployment in deployments + ), provider_live=True, ) @@ -632,18 +639,14 @@ def _compat_models_registered() -> Any: yield return + from proxy_client import ModelNotRegistered, ModelRegistered from requests import RequestException proxy = _build_control_plane_client(proxy_config) - registered_ids: list[str] = [] - failures: list[tuple[str, str]] = [] + registrations = _register_deployments(proxy, load_all_deployments()) + registered_ids = tuple(r.model_id for r in registrations if isinstance(r, ModelRegistered)) + failures = tuple((r.model_name, r.reason) for r in registrations if isinstance(r, ModelNotRegistered)) try: - for deployment in load_all_deployments(): - try: - model_id = _register_deployment(proxy, deployment) - registered_ids.append(model_id) - except (AssertionError, RequestException) as exc: - failures.append((deployment.model_name, str(exc))) if failures: summary = "\n".join( f" - {name}: {reason}" for name, reason in failures diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 3bb4ab0d3be..309168351e7 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -37,6 +37,7 @@ from e2e_config import ( PROMPT_CACHING_OPT_IN_ENV, PROVIDER_EDGE_HOST_OPT_IN_ENV, PROXY_BASE_URL, + PROXY_WIDE_CONFIG_OPT_IN_ENV, SAGEMAKER_NOVA2_LITE_OPT_IN_ENV, SAGEMAKER_NOVA_OPT_IN_ENV, SECRET_MANAGER_OPT_IN_ENV, @@ -56,7 +57,6 @@ from models import TeamNewBody, UserNewBody, UserNewResponse from provider_cache_routing import LIVE_PROVIDER_REQUIRED from provider_edge import replay_leftover_error from proxy_client import ProxyClient, build_proxy_client -from stack_lock import stack_lock _E2E_TEST_RAN = pytest.StashKey[bool]() _CALL_PASSED = pytest.StashKey[bool]() @@ -80,6 +80,7 @@ OPT_IN_MARKERS: Final = MappingProxyType( "secret_manager": SECRET_MANAGER_OPT_IN_ENV, "sagemaker_nova": SAGEMAKER_NOVA_OPT_IN_ENV, "sagemaker_nova2_lite": SAGEMAKER_NOVA2_LITE_OPT_IN_ENV, + "proxy_wide_config": PROXY_WIDE_CONFIG_OPT_IN_ENV, } ) @@ -165,11 +166,6 @@ def pytest_configure(config: pytest.Config) -> None: "redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from " "gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set", ) - config.addinivalue_line( - "markers", - "quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; " - "every other test waits for it to finish", - ) config.addinivalue_line( "markers", "mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless " @@ -208,6 +204,11 @@ def pytest_configure(config: pytest.Config) -> None: "sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; " "deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it", ) + config.addinivalue_line( + "markers", + "proxy_wide_config: turns on a proxy-wide setting through /config/update, so it runs on a stack no other suite " + "is driving; deselected unless E2E_PROXY_WIDE_CONFIG_STACK is set", + ) def pytest_sessionstart(session: pytest.Session) -> None: @@ -277,7 +278,7 @@ def pytest_collection_finish(session: pytest.Session) -> None: ahead of xdist's own collection-finish report, and the controller schedules no test until every worker has reported, so this is the idle footprint of a stack that just passed its readiness gate. The fixture hands the capture to the - idle-budget test in router/test_reliability_memory_e2e.py.""" + idle-budget test in perf/test_reliability_memory_e2e.py.""" if session.config.getoption("collectonly") or not any(_uses_idle_rss(item) for item in session.items): return session.config.stash[_IDLE_RSS] = read_rss_everywhere(build_proxy_client(), timeout=IDLE_RSS_READ_TIMEOUT_SECONDS) @@ -307,12 +308,6 @@ def _proxy_fail_reason() -> str | None: return None -@pytest.hookimpl(wrapper=True) -def pytest_runtest_protocol(item: pytest.Item, nextitem: pytest.Item | None) -> Generator[None, object, object]: - with stack_lock(exclusive=item.get_closest_marker("quiet_stack") is not None): - return (yield) - - @pytest.hookimpl(tryfirst=True) def pytest_runtest_setup(item: pytest.Item) -> None: """Hard-fail `e2e`-marked tests unless a proxy answers its liveness probe. diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index df554b20204..c44e0d663bf 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -44,8 +44,6 @@ - {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/pass_through/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"} - {id: quota_management.spend_tracking.embeddings.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: embeddings, assertions: [logs_cost], exercised_on: [embeddings], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Embedding calls write nonzero spend rows"} - {id: quota_management.spend_tracking.cache_hit.zero_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cache_hit, assertions: [zero_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A response-cache hit logs at zero cost with the cache-hit marker"} -- {id: quota_management.spend_tracking.discount_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: discount_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider discount changes the real request's spend row"} -- {id: quota_management.spend_tracking.margin_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: margin_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider margin is added to the real request's spend row"} - {id: quota_management.spend_tracking.guardrail_cost.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: guardrail_cost, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/guardrail_cost.py + proxy/spend_tracking/spend_tracking_utils.py", rationale: "A priced Bedrock guardrail charge reaches the persisted spend row"} - {id: quota_management.spend_tracking.key_rollup.matches_sum_of_logs, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_rollup, assertions: [matches_sum_of_logs], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A key's rolled-up spend equals the sum of its log rows"} - {id: quota_management.spend_tracking.concurrent_burst.loses_no_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: concurrent_burst, assertions: [loses_no_spend], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "Concurrent calls all land as spend; no row lost to write contention"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index efbc95966c4..bcab7e35ae7 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -215,6 +215,7 @@ OTEL_TLS_OPT_IN_ENV: Final = "E2E_OTEL_EXPORTER_ENDPOINT" SECRET_MANAGER_OPT_IN_ENV: Final = "E2E_SECRET_MANAGER" SAGEMAKER_NOVA_OPT_IN_ENV: Final = "SAGEMAKER_NOVA_ENDPOINT" SAGEMAKER_NOVA2_LITE_OPT_IN_ENV: Final = "SAGEMAKER_NOVA2_LITE_ENDPOINT" +PROXY_WIDE_CONFIG_OPT_IN_ENV: Final = "E2E_PROXY_WIDE_CONFIG_STACK" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) @@ -239,7 +240,6 @@ MEMORY_RSS_SETTLE_SAMPLES = int(os.environ.get("E2E_MEMORY_RSS_SETTLE_SAMPLES", MEMORY_RSS_SAMPLE_INTERVAL_SECONDS = float(os.environ.get("E2E_MEMORY_RSS_SAMPLE_INTERVAL_SECONDS", "1")) MEMORY_RSS_CEILING_MB = float(os.environ.get("E2E_MEMORY_RSS_CEILING_MB", "200")) MEMORY_IDLE_RSS_BUDGET_MB = float(os.environ.get("E2E_MEMORY_IDLE_RSS_BUDGET_MB", "768")) -MEMORY_STORED_REQUEST_BUDGET_KB = float(os.environ.get("E2E_MEMORY_STORED_REQUEST_BUDGET_KB", "64")) def ws_base_url() -> str: diff --git a/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py b/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py index 6e564f4625a..1ec2a6822e3 100644 --- a/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py +++ b/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py @@ -5,7 +5,8 @@ embedding deployment, and turns on `litellm_settings.mcp_semantic_tool_filter` t /config/update. A chat completion that attaches the server and asks to search logs must report, in the x-litellm-semantic-filter headers, that the catalog was cut to at most top_k tools, all of them tools the key lists. Requires DD_API_KEY + DD_APP_KEY and the proxy's -OPENAI_API_KEY for the embedding deployment. +OPENAI_API_KEY for the embedding deployment. The filter is proxy-wide, so the test is +marked `proxy_wide_config` and runs on a stack the main suite is not driving. """ from __future__ import annotations @@ -23,7 +24,7 @@ from mcp_client import McpClient from models import ChatBody, ChatMessage, LiteLLMParamsBody, McpChatTool from pydantic import BaseModel -pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack] +pytestmark = [pytest.mark.e2e, pytest.mark.proxy_wide_config] TOP_K: Final = 3 diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 662b307dfe8..d42518650fb 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1295,27 +1295,6 @@ class CostMap(RootModel[dict[str, CostMapEntry]]): pass -class ConfigPatchResponse(BaseModel): - status: str - values: dict[str, float | dict[str, float]] - - -class CostDiscountConfig(RootModel[dict[str, float]]): - pass - - -class CostMarginConfig(RootModel[dict[str, float | dict[str, float]]]): - pass - - -class CostDiscountConfigResponse(BaseModel): - values: dict[str, float] - - -class CostMarginConfigResponse(BaseModel): - values: dict[str, float | dict[str, float]] - - class BedrockGuardrailParams(BaseModel): guardrail: Literal["bedrock"] = "bedrock" mode: Literal["pre_call"] = "pre_call" diff --git a/tests/e2e/perf/test_perf_cases.py b/tests/e2e/perf/test_perf_cases.py index d89a8634ab3..167c743d8b7 100644 --- a/tests/e2e/perf/test_perf_cases.py +++ b/tests/e2e/perf/test_perf_cases.py @@ -32,7 +32,7 @@ from setup.podcpu import gateway_cpu_sampler from setup.provision import MOCK_PROVIDERS, build_setup, guardrail_url, mock_url from setup.run import run_plan -pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack] +pytestmark = pytest.mark.perf _SPEND_SAMPLE: Final = 40 diff --git a/tests/e2e/router/test_reliability_memory_e2e.py b/tests/e2e/perf/test_reliability_memory_e2e.py similarity index 77% rename from tests/e2e/router/test_reliability_memory_e2e.py rename to tests/e2e/perf/test_reliability_memory_e2e.py index fb114092303..2066137bb72 100644 --- a/tests/e2e/router/test_reliability_memory_e2e.py +++ b/tests/e2e/perf/test_reliability_memory_e2e.py @@ -1,17 +1,19 @@ """Live e2e: a few hundred requests that fail before any provider answers must not -grow the proxy's resident memory past a loose ceiling once the proxy is warm. +grow the proxy's resident memory past a loose ceiling once the proxy is warm. It reads +the proxy's own RSS, so it runs in the perf lane, where nothing else is hitting the stack. The regression this guards shipped in v1.100.0: every retry breadcrumb copied the whole request and the copies nested into one router-global list, so a proxy under retry-heavy failing traffic grew until it was OOM-killed. The traffic here has that shape: a model group whose deployments refuse at the socket (an unreachable base -URL) with cooldown_time 0 so the router keeps retrying them, per-request retries, -and a fallback group that refuses the same way, each request carrying a long chat +URL, with a literal key since the perf stack carries no provider keys) with +cooldown_time 0 so the router keeps retrying them, per-request retries, and a +fallback group that refuses the same way, each request carrying a long chat transcript so every whole-request copy costs hundreds of containers instead of a -handful. Under the stack's cooldown policy a -deployment that fails a handful of times in a row is benched (a bad-credential 401 -included), the router answers "No deployments available" without retrying, and the -retry loop that leaks stops running; cooldown_time 0 keeps it running. +handful. Under the stack's cooldown policy a deployment that fails a handful of +times in a row is benched (a bad-credential 401 included), the router answers "No +deployments available" without retrying, and the retry loop that leaks stops running; +cooldown_time 0 keeps it running. Four identical batches run back to back. The first two are the warmup that grows the proxy's caches and allocator arenas, the last two are the measured window, and @@ -35,19 +37,11 @@ single load-balanced address reaches the workers behind it one answer at a time, a worker that answered only one checkpoint would otherwise drop out of the comparison, which is where a leaking worker could hide. -RSS alone is a coarse gauge: on the release stack (spend logs storing prompts, -json logs, prometheus and otel callbacks) the same v1.100.0 breadcrumbs grew RSS -by only about 15 MB per 300 failing requests, while every failing request's stored -request snapshot carried a copy of the request per failed attempt, over 100 KB on -the first call and a couple of MB once the copies nested, against tens of KB with -the fix. So the first check sends one failing request before the phases, reads its -spend log back through /spend/logs, and holds the stored request body to a fixed -size budget: the deterministic catch for a breadcrumb that copies the whole -request. It runs before the phases because the leaking writer drops its own rows -under the phases' traffic (a queue budget hit, a recursion limit on the nested -copies), which would turn the size check into a missing-row check. +The stored-request half of that regression (every failing request's spend log carrying a +copy of the request per failed attempt) is a deterministic size check, so it lives in +tests/integration/spend/test_failing_request_stored_snapshot.py. -A second, cheaper check holds the idle footprint: every worker's RSS as the harness +A cheaper check holds the idle footprint: every worker's RSS as the harness read it at collection time, before this pytest process sent any traffic (see conftest.pytest_collection_finish), must sit under a fixed budget. On the release gate that is a fresh stack right after its readiness gate, one worker per @@ -74,7 +68,6 @@ from typing import Final import pytest -from complexity_router_client import ComplexityRouterClient from e2e_config import ( MEMORY_CONCURRENCY, MEMORY_IDLE_RSS_BUDGET_MB, @@ -83,7 +76,6 @@ from e2e_config import ( MEMORY_RSS_CEILING_MB, MEMORY_RSS_SAMPLE_INTERVAL_SECONDS, MEMORY_RSS_SETTLE_SAMPLES, - MEMORY_STORED_REQUEST_BUDGET_KB, MEMORY_TRANSCRIPT_TURNS, PROXY_REPLICA_URLS, unique_marker, @@ -91,11 +83,12 @@ from e2e_config import ( from e2e_metadata import Domain, Mode, Provider, Subject, meta from lifecycle import ResourceManager from memory_readings import RssCapture, RssReading, WorkerKey, read_rss_everywhere -from models import ChatMessage, RouterSettingsOverride, SpendLogRow +from models import ChatMessage, LiteLLMParamsBody, ReliabilityChatBody, RouterSettingsOverride from proxy_client import ProxyClient -from reliability_support import REAL_MODEL, chat_override, create_never_benched_refusing_deployment -pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack] +pytestmark = [pytest.mark.e2e, pytest.mark.perf] + +REFUSING_MODEL: Final = "openai/gpt-5.5" DEPLOYMENTS_PER_GROUP: Final = 2 RSS_SAMPLE_CAP: Final = 4 * MEMORY_RSS_SETTLE_SAMPLES @@ -110,7 +103,6 @@ class FailedCall: status_code: int seconds: float body_head: str - call_id: str | None @dataclass(frozen=True, slots=True) @@ -123,8 +115,17 @@ class WorkerGrowth: return self.after.ram_usage_mb - self.warm.ram_usage_mb +def _create_never_benched_refusing_deployment(proxy: ProxyClient, name: str) -> str: + return proxy.create_model( + name, + LiteLLMParamsBody( + model=REFUSING_MODEL, api_key="sk-unreachable", api_base="http://127.0.0.1:9/v1", cooldown_time=0 + ), + ) + + def _register_refusing_group(proxy: ProxyClient, resources: ResourceManager, name: str) -> None: - for model_id in tuple(create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)): + for model_id in tuple(_create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)): resources.defer(lambda model_id=model_id: proxy.delete_model(model_id)) @@ -141,10 +142,18 @@ TRANSCRIPT: Final = _transcript(MEMORY_TRANSCRIPT_TURNS) def _fail_once(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> FailedCall: started: Final = time.perf_counter() - resp: Final = chat_override( - proxy, key, model, f"memory regression {unique_marker()}", override=override, history=TRANSCRIPT + resp: Final = proxy.transport.send( + "/chat/completions", + headers=proxy.transport.bearer(key), + json=ReliabilityChatBody( + model=model, + messages=[*TRANSCRIPT, ChatMessage(role="user", content=f"memory regression {unique_marker()}")], + max_tokens=512, + router_settings_override=override, + cache={"no-cache": True}, + ), ) - return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300], resp.call_id) + return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300]) def _fail_many(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> tuple[FailedCall, ...]: @@ -259,24 +268,6 @@ def _assert_every_call_failed_through_fallback(calls: Sequence[FailedCall], fall ) -def _stored_request_kb(proxy: ProxyClient, call: FailedCall) -> float: - assert call.call_id, ( - f"the failing call carried no x-litellm-call-id header, so its spend log cannot be read back: {call.body_head}" - ) - rows: Final[Sequence[SpendLogRow]] = proxy.poll_logs_for_request_id(call.call_id) - assert rows, ( - f"no spend log row appeared for failing call {call.call_id} within the poll window: either the stack " - "writes no spend logs or its writer dropped the row, which the v1.100.0 one did once the stored " - "request outgrew the writer's queue budget" - ) - snapshot: Final = rows[0].proxy_server_request - assert snapshot, ( - f"spend log {call.call_id} stored no request body, so the stack is not running with " - "general_settings.store_prompts_in_spend_logs and the stored-request check would pass vacuously" - ) - return len(json.dumps(snapshot).encode()) / 1024 - - class TestReliabilityMemory: @pytest.mark.covers("reliability.perf.idle_memory.under_slo") @meta( @@ -309,13 +300,13 @@ class TestReliabilityMemory: Subject( domain=Domain.ROUTING, providers=(Provider.OPENAI,), - models=(REAL_MODEL,), + models=(REFUSING_MODEL,), mode=Mode.NONSTREAM, ) ) - def test_failing_requests_do_not_grow_rss_or_stored_request( + def test_failing_requests_do_not_grow_rss( self, - client: ComplexityRouterClient, + proxy: ProxyClient, resources: ResourceManager, scoped_key: str, record_property: Callable[[str, object], None], @@ -323,35 +314,25 @@ class TestReliabilityMemory: marker: Final = unique_marker() primary: Final = f"reliability-memory-{marker}" fallback: Final = f"reliability-memory-fb-{marker}" - _register_refusing_group(client.proxy, resources, primary) - _register_refusing_group(client.proxy, resources, fallback) + _register_refusing_group(proxy, resources, primary) + _register_refusing_group(proxy, resources, fallback) override: Final = RouterSettingsOverride( num_retries=MEMORY_RETRIES_PER_REQUEST, fallbacks=[{primary: [fallback]}] ) - probe: Final = _fail_once(client.proxy, scoped_key, primary, override) - _assert_every_call_failed_through_fallback((probe,), fallback) - stored_kb: Final = _stored_request_kb(client.proxy, probe) - assert stored_kb <= MEMORY_STORED_REQUEST_BUDGET_KB, ( - f"the spend log of one failing request stored a {stored_kb:.0f} KB request body, past the " - f"{MEMORY_STORED_REQUEST_BUDGET_KB:.0f} KB budget for a {len(TRANSCRIPT)}-message transcript with " - f"{MEMORY_RETRIES_PER_REQUEST} retries and a fallback; the retry breadcrumbs are copying the whole " - f"request into the stored snapshot the way the v1.100.0 ones did" - ) - warm: Final = _settled_rss_after_batches( - client.proxy, + proxy, scoped_key, primary, override, fallback, record_property, - _settled_rss_per_worker(client.proxy), + _settled_rss_per_worker(proxy), 1, WARMUP_BATCHES, ) after: Final = _settled_rss_after_batches( - client.proxy, + proxy, scoped_key, primary, override, diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 8cc7bf0c2b4..9b842a723d8 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -11,7 +11,7 @@ from __future__ import annotations import os import time import warnings -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass, field, replace from datetime import datetime from functools import reduce @@ -155,6 +155,21 @@ class NotServableOn: last_result: Result[ModelsListResponse] | None +@dataclass(frozen=True, slots=True) +class ModelRegistered: + model_name: str + model_id: str + + +@dataclass(frozen=True, slots=True) +class ModelNotRegistered: + model_name: str + reason: str + + +ModelRegistration = ModelRegistered | ModelNotRegistered + + def await_servable( list_models: Callable[[float], Result[ModelsListResponse]], *, @@ -461,6 +476,7 @@ class ProxyClient: model_servable_db_sync_seconds: float = MODEL_SERVABLE_DB_SYNC_SECONDS model_servable_interval: float = MODEL_SERVABLE_INTERVAL model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT + settle: Callable[[float], None] = settle_propagation def with_caller(self, caller: Caller) -> ProxyClient: return replace(self, caller=caller) @@ -694,31 +710,59 @@ class ProxyClient: balancer address is configured (every request opens a fresh connection, so the caller's next request re-rolls), so waiting out PROPAGATION_TIMEOUT is what makes the model safe to use anywhere.""" - model_id = unwrap( - self.transport.post( - "/model/new", - headers=self.management_headers(), - json=body.model_copy(update={"litellm_params": route_cache_model( - body.litellm_params, provider_edge_base, - enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1" and not provider_live, - mode=body.model_info.mode, - )}), - response_type=ModelNewResponse, - ) - ).model_id + model_id = unwrap(self._post_model(body, provider_live=provider_live)).model_id written_at = time.monotonic() try: self._await_model_servable(body.model_name, listed_for) except BaseException: self.delete_model(model_id) raise - settle_propagation(written_at) + self.settle(written_at) return model_id + @step("Add a batch of deployments, then wait once for all of them to propagate") + def register_models( + self, bodies: Sequence[ModelNewBody], *, provider_live: bool = False + ) -> tuple[ModelRegistration, ...]: + posted: Final = tuple((body.model_name, self._post_model(body, provider_live=provider_live)) for body in bodies) + written_at: Final = time.monotonic() + registrations: Final = tuple(self._servable_registration(name, result) for name, result in posted) + self.settle(written_at) + return registrations + + def _post_model(self, body: ModelNewBody, *, provider_live: bool) -> Result[ModelNewResponse]: + return self.transport.post( + "/model/new", + headers=self.management_headers(), + json=body.model_copy(update={"litellm_params": route_cache_model( + body.litellm_params, provider_edge_base, + enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1" and not provider_live, + mode=body.model_info.mode, + )}), + response_type=ModelNewResponse, + ) + + def _servable_registration(self, model_name: str, posted: Result[ModelNewResponse]) -> ModelRegistration: + if not isinstance(posted, Success): + return ModelNotRegistered(model_name=model_name, reason=f"/model/new failed: {posted!r}") + match self._servable_outcome(model_name, None): + case Servable(): + return ModelRegistered(model_name=model_name, model_id=posted.data.model_id) + case NotServableOn() as missing: + self.delete_model(posted.data.model_id) + return ModelNotRegistered(model_name=model_name, reason=self._not_servable_message(model_name, missing)) + def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None: """Block until every replica lists `model_name`, or fail at model_servable_timeout.""" + match self._servable_outcome(model_name, listed_for): + case Servable(): + return + case NotServableOn() as missing: + raise AssertionError(self._not_servable_message(model_name, missing)) + + def _servable_outcome(self, model_name: str, listed_for: str | None) -> Servable | NotServableOn: headers: Final = self.management_headers(listed_for) - outcome: Final = await_servable_everywhere( + return await_servable_everywhere( {url: self._models_poller(transport, headers) for url, transport in self.replicas.items()}, model_name=model_name, timeout=self.model_servable_timeout, @@ -728,19 +772,15 @@ class ProxyClient: now=time.monotonic, sleep=time.sleep, ) - match outcome: - case Servable(): - return - case NotServableOn(replica=replica, last_result=last_result): - raise AssertionError( - servable_timeout_message( - model_name=model_name, - replica=replica, - timeout=self.model_servable_timeout, - db_sync_seconds=self.model_servable_db_sync_seconds, - last_result=last_result, - ) - ) + + def _not_servable_message(self, model_name: str, missing: NotServableOn) -> str: + return servable_timeout_message( + model_name=model_name, + replica=missing.replica, + timeout=self.model_servable_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, + last_result=missing.last_result, + ) @staticmethod def _models_poller(transport: Transport, headers: AuthHeaders) -> ModelsPoller: diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index da3dd72d795..eab59c3c77c 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -13,7 +13,6 @@ markers = prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set - quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; every other test waits for it to finish mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless E2E_MCP_OAUTH_LIVE is set provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set owned_gateway: boots its own proxy from source against the stack's Postgres, so it needs DATABASE_URL on the pytest host; deselected unless E2E_OWNED_GATEWAY is set @@ -22,3 +21,4 @@ markers = secret_manager: needs a proxy booted from gateway/secret_manager__ci_config.yml against that live secret manager; deselected unless E2E_SECRET_MANAGER names the backend (see secret_manager/secret_backends.py) sagemaker_nova: needs a deployed SageMaker Nova endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA_ENDPOINT names it sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it + proxy_wide_config: turns on a proxy-wide setting through /config/update, so it runs on a stack no other suite is driving; deselected unless E2E_PROXY_WIDE_CONFIG_STACK is set diff --git a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py b/tests/e2e/quota_management/spend_tracking/test_guardrail_cost_spend_e2e.py similarity index 52% rename from tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py rename to tests/e2e/quota_management/spend_tracking/test_guardrail_cost_spend_e2e.py index 8f3bed3f89b..a77ca91ad3a 100644 --- a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_guardrail_cost_spend_e2e.py @@ -17,11 +17,6 @@ from models import ( ChatBody, ChatMessage, ChatResponse, - ConfigPatchResponse, - CostDiscountConfig, - CostDiscountConfigResponse, - CostMarginConfig, - CostMarginConfigResponse, GuardrailCreateBody, GuardrailCreateResponse, KeyDeleteBody, @@ -31,14 +26,11 @@ from models import ( ) from spend_e2e_client import SpendClient -pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack] +pytestmark = pytest.mark.e2e BACKEND_MODEL: Final = "openai/gpt-6.1-sol" INPUT_RATE: Final = 0.00005 OUTPUT_RATE: Final = 0.0001 -DISCOUNT: Final = 0.25 -MARGIN_PERCENT: Final = 0.1 -MARGIN_FIXED: Final = 0.0005 def _register_bedrock_guardrail( @@ -92,52 +84,6 @@ def _guarded_chat(client: SpendClient, key: str, model: str, name: str) -> Strea ) -def _set_discount(client: SpendClient, values: dict[str, float]) -> float: - unwrap( - client.proxy.transport.patch( - "/config/cost_discount_config", - headers=client.proxy.transport.master, - json=CostDiscountConfig(values), - response_type=ConfigPatchResponse, - ) - ) - return time.monotonic() - - -def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]]) -> float: - unwrap( - client.proxy.transport.patch( - "/config/cost_margin_config", - headers=client.proxy.transport.master, - json=CostMarginConfig(values), - response_type=ConfigPatchResponse, - ) - ) - return time.monotonic() - - -def _get_discount(client: SpendClient) -> dict[str, float]: - return unwrap( - client.proxy.transport.get( - "/config/cost_discount_config", - headers=client.proxy.transport.master, - params=NoBody(), - response_type=CostDiscountConfigResponse, - ) - ).values - - -def _get_margin(client: SpendClient) -> dict[str, float | dict[str, float]]: - return unwrap( - client.proxy.transport.get( - "/config/cost_margin_config", - headers=client.proxy.transport.master, - params=NoBody(), - response_type=CostMarginConfigResponse, - ) - ).values - - def _register_model(client: SpendClient, resources: ResourceManager, prefix: str) -> str: model: Final = f"{prefix}-{unique_marker()}" model_id: Final = client.proxy.create_model( @@ -199,100 +145,7 @@ def scoped_key(client: SpendClient, strict_resources: ResourceManager) -> str: return key -@pytest.fixture -def restored_pricing_config(client: SpendClient) -> Iterator[None]: - discount: Final = _get_discount(client) - margin: Final = _get_margin(client) - try: - _set_discount(client, {}) - reset_at: Final = _set_margin(client, {}) - settle_propagation(reset_at) - yield - finally: - try: - _set_discount(client, discount) - finally: - restored_at: Final = _set_margin(client, margin) - settle_propagation(restored_at) - - -class TestPricingConfigSpend: - @pytest.mark.covers( - "quota_management.spend_tracking.discount_config.logs_cost", - exercised_on=["chat_completions"], - ) - @meta( - Subject( - domain=Domain.SPEND_BUDGETS, - route=Route.SPEND_REPORTING, - providers=(Provider.OPENAI,), - models=(BACKEND_MODEL,), - mode=Mode.NONSTREAM, - ) - ) - def test_configured_discount_reaches_persisted_spend_row( - self, - client: SpendClient, - strict_resources: ResourceManager, - scoped_key: str, - restored_pricing_config: None, - ) -> None: - discount_set_at: Final = _set_discount(client, {"openai": DISCOUNT}) - settle_propagation(discount_set_at) - model: Final = _register_model(client, strict_resources, "discount-priced") - chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16)) - assert chat.id and chat.usage - assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0 - assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0 - - row: Final = poll_cost_row(client.proxy, chat.id) - assert row is not None - base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens) - breakdown: Final = row.breakdown - assert breakdown.discount_percent is not None and approx_equal(breakdown.discount_percent, DISCOUNT) - assert breakdown.discount_amount is not None and approx_equal(breakdown.discount_amount, base_cost * DISCOUNT) - assert row.spend is not None and approx_equal(row.spend, base_cost * (1 - DISCOUNT)) - - @pytest.mark.covers( - "quota_management.spend_tracking.margin_config.logs_cost", - exercised_on=["chat_completions"], - ) - @meta( - Subject( - domain=Domain.SPEND_BUDGETS, - route=Route.SPEND_REPORTING, - providers=(Provider.OPENAI,), - models=(BACKEND_MODEL,), - mode=Mode.NONSTREAM, - ) - ) - def test_configured_margin_reaches_persisted_spend_row( - self, - client: SpendClient, - strict_resources: ResourceManager, - scoped_key: str, - restored_pricing_config: None, - ) -> None: - margin_set_at: Final = _set_margin( - client, {"openai": {"percentage": MARGIN_PERCENT, "fixed_amount": MARGIN_FIXED}} - ) - settle_propagation(margin_set_at) - model: Final = _register_model(client, strict_resources, "margin-priced") - chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16)) - assert chat.id and chat.usage - assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0 - assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0 - - row: Final = poll_cost_row(client.proxy, chat.id) - assert row is not None - base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens) - breakdown: Final = row.breakdown - expected_margin: Final = base_cost * MARGIN_PERCENT + MARGIN_FIXED - assert breakdown.margin_percent is not None and approx_equal(breakdown.margin_percent, MARGIN_PERCENT) - assert breakdown.margin_fixed_amount is not None and approx_equal(breakdown.margin_fixed_amount, MARGIN_FIXED) - assert breakdown.margin_total_amount is not None and approx_equal(breakdown.margin_total_amount, expected_margin) - assert row.spend is not None and approx_equal(row.spend, base_cost + expected_margin) - +class TestGuardrailCostSpend: @pytest.mark.covers( "quota_management.spend_tracking.guardrail_cost.logs_cost", exercised_on=["chat_completions"], @@ -311,7 +164,6 @@ class TestPricingConfigSpend: client: SpendClient, strict_resources: ResourceManager, scoped_key: str, - restored_pricing_config: None, ) -> None: identifier: Final = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"] version: Final = os.environ["BEDROCK_GUARDRAIL_VERSION"] diff --git a/tests/e2e/stack_lock.py b/tests/e2e/stack_lock.py deleted file mode 100644 index 06df7a20b6a..00000000000 --- a/tests/e2e/stack_lock.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Cross-process reader/writer lock over the proxy stack every xdist worker shares. -Every collected test holds it shared, marker or not, since the Claude Code cells and -other unmarked suites drive the same stack; a `quiet_stack` test holds it exclusive, -and the `gate` file makes a waiting exclusive holder win over readers that arrive -after it.""" - -from __future__ import annotations - -import fcntl -import hashlib -import tempfile -from collections.abc import Generator -from contextlib import ExitStack, contextmanager -from pathlib import Path -from typing import Final - -from e2e_config import PROXY_BASE_URL - -STACK_DIGEST: Final = hashlib.sha256(PROXY_BASE_URL.encode()).hexdigest()[:12] -LOCK_DIR: Final = Path(tempfile.gettempdir()) / f"litellm-e2e-stack-{STACK_DIGEST}" -GATE_FILE: Final = LOCK_DIR / "gate" -STACK_FILE: Final = LOCK_DIR / "stack" - - -@contextmanager -def _flock(path: Path, operation: int) -> Generator[None]: - with path.open("a") as handle: - fcntl.flock(handle, operation) - try: - yield - finally: - fcntl.flock(handle, fcntl.LOCK_UN) - - -@contextmanager -def stack_lock(exclusive: bool) -> Generator[None]: - LOCK_DIR.mkdir(parents=True, exist_ok=True) - if exclusive: - with _flock(GATE_FILE, fcntl.LOCK_EX), _flock(STACK_FILE, fcntl.LOCK_EX): - yield - return - with ExitStack() as held: - with _flock(GATE_FILE, fcntl.LOCK_SH): - held.enter_context(_flock(STACK_FILE, fcntl.LOCK_SH)) - yield diff --git a/tests/e2e_harness/AGENTS.md b/tests/e2e_harness/AGENTS.md index 65a3b32df96..54eee2fa5ed 100644 --- a/tests/e2e_harness/AGENTS.md +++ b/tests/e2e_harness/AGENTS.md @@ -1,6 +1,6 @@ # e2e harness tests -Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the stack lock, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network +Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network The layout mirrors `tests/e2e/`: `test_e2e_http.py` covers `tests/e2e/e2e_http.py`, `logging/test_datadog_reader.py` covers `tests/e2e/logging/datadog_reader.py`, and `claude_code/` covers the driver, builder, probe and version resolver. Put a new harness test under the folder that mirrors the suite folder whose module it covers diff --git a/tests/e2e_harness/test_proxy_client.py b/tests/e2e_harness/test_proxy_client.py index e1615ec5f65..d352e5c8d36 100644 --- a/tests/e2e_harness/test_proxy_client.py +++ b/tests/e2e_harness/test_proxy_client.py @@ -12,10 +12,11 @@ injected clock, so nothing here monkeypatches anything. from __future__ import annotations import json +import time from builtins import ExceptionGroup from collections.abc import Callable, Generator, Iterable, Mapping from contextlib import contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, replace from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from itertools import chain, repeat from queue import SimpleQueue @@ -40,7 +41,10 @@ from models import ( LiteLLMParamsBody, McpServerCreateBody, McpServerUpdateBody, + ModelDeleteBody, + ModelInfoBody, ModelListEntry, + ModelNewBody, ModelsListResponse, OrgNewBody, OrgUpdateBody, @@ -59,6 +63,8 @@ from proxy_client import ( ConvergeOutcome, CredentialKind, EverywhereConverged, + ModelNotRegistered, + ModelRegistered, ModelsPoller, NeverConvergedOn, NotConverged, @@ -691,3 +697,85 @@ def test_partial_updates_preserve_explicit_null_at_the_http_boundary(operation: ) assert json.loads(bodies.get_nowait()) == expected assert bodies.empty() + + +@contextmanager +def model_registry( + listed: frozenset[str], refused: frozenset[str] +) -> Generator[tuple[ProxyClient, SimpleQueue[float], SimpleQueue[str]]]: + writes: Final[SimpleQueue[float]] = SimpleQueue() + deleted: Final[SimpleQueue[str]] = SimpleQueue() + + class Handler(BaseHTTPRequestHandler): + def log_message(self, format: str, *args: object) -> None: + pass + + def reply(self, status: int, body: object) -> None: + self.send_response(status) + self.end_headers() + self.wfile.write(json.dumps(body).encode()) + + def do_GET(self) -> None: + self.reply(200, {"data": [{"id": name} for name in sorted(listed)]}) + + def do_POST(self) -> None: + raw: Final = self.rfile.read(int(self.headers.get("Content-Length", "0"))) + if self.path == "/model/delete": + deleted.put(ModelDeleteBody.model_validate_json(raw).id) + self.reply(200, {}) + return + writes.put(time.monotonic()) + name: Final = ModelNewBody.model_validate_json(raw).model_name + if name in refused: + self.reply(500, {"error": "refused"}) + return + self.reply(200, {"model_id": f"id-{name}"}) + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) + thread.start() + url: Final = f"http://127.0.0.1:{server.server_port}" + proxy: Final = build_proxy_client( + base_url=url, control_plane_base_url=url, replica_urls=(url,), control_replica_urls=(url,), master_key="bootstrap" + ) + try: + yield replace(proxy, model_servable_timeout=0.3, model_servable_interval=0.01), writes, deleted + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +class TestRegisterModels: + def test_a_batch_waits_once_after_its_last_write_and_reports_each_model(self) -> None: + names: Final = ("first", "refused", "ghost", "last") + settled: Final[SimpleQueue[float]] = SimpleQueue() + with ( + model_registry(listed=frozenset({"first", "last"}), refused=frozenset({"refused"})) as ( + proxy, + writes, + deleted, + ), + without_retries(), + ): + registrations: Final = replace(proxy, settle=settled.put).register_models( + tuple( + ModelNewBody( + model_name=name, litellm_params=LiteLLMParamsBody(model="synthetic"), model_info=ModelInfoBody() + ) + for name in names + ) + ) + last_write: Final = max(writes.get_nowait() for _ in names) + assert writes.empty() + assert settled.get_nowait() >= last_write + assert settled.empty(), "a batch must wait out propagation once, not once per model" + assert tuple((type(r), r.model_name) for r in registrations) == ( + (ModelRegistered, "first"), + (ModelNotRegistered, "refused"), + (ModelNotRegistered, "ghost"), + (ModelRegistered, "last"), + ) + assert tuple(r.model_id for r in registrations if isinstance(r, ModelRegistered)) == ("id-first", "id-last") + assert deleted.get_nowait() == "id-ghost" + assert deleted.empty() diff --git a/tests/e2e_harness/test_stack_lock.py b/tests/e2e_harness/test_stack_lock.py deleted file mode 100644 index b275ca8c755..00000000000 --- a/tests/e2e_harness/test_stack_lock.py +++ /dev/null @@ -1,117 +0,0 @@ -"""Cross-process behavior of the stack lock: readers share it, an exclusive holder waits for -every reader and keeps them out, and a reader arriving behind a waiting exclusive holder -queues behind it instead of starving it.""" - -from __future__ import annotations - -import fcntl -import inspect -import os -import subprocess -import sys -import time -from contextlib import ExitStack -from pathlib import Path -from typing import Final - -import pytest -import stack_lock - -HARNESS_DIR: Final = Path(inspect.getfile(stack_lock)).resolve().parent -DEADLINE_SECONDS: Final = 30.0 -SETTLE_SECONDS: Final = 0.5 -HOLDER_SCRIPT: Final = """ -import sys, time -from pathlib import Path -from stack_lock import stack_lock -name, mode, release_path, log_path = sys.argv[1:] - - -def record(event): - with Path(log_path).open("a") as log: - log.write(f"{name} {event}\\n") - - -record("waiting") -with stack_lock(exclusive=mode == "exclusive"): - record("enter") - while not Path(release_path).exists(): - time.sleep(0.02) - record("exit") -""" - - -def _events(log_path: Path) -> tuple[str, ...]: - return tuple(log_path.read_text().splitlines()) if log_path.exists() else () - - -def _wait_for_event(log_path: Path, event: str) -> None: - deadline: Final = time.monotonic() + DEADLINE_SECONDS - while event not in _events(log_path): - if time.monotonic() > deadline: - pytest.fail(f"{event!r} never appeared; events so far: {_events(log_path)}") - time.sleep(0.02) - - -def _wait_until_gate_is_held_exclusively(gate_path: Path) -> None: - deadline: Final = time.monotonic() + DEADLINE_SECONDS - with gate_path.open("a") as handle: - while True: - try: - fcntl.flock(handle, fcntl.LOCK_SH | fcntl.LOCK_NB) - except BlockingIOError: - return - fcntl.flock(handle, fcntl.LOCK_UN) - if time.monotonic() > deadline: - pytest.fail("no exclusive holder ever took the gate") - time.sleep(0.02) - - -def _start_holder(held: ExitStack, tmp_path: Path, name: str, mode: str) -> subprocess.Popen[bytes]: - holder: Final = held.enter_context( - subprocess.Popen( - ( - sys.executable, - "-P", - "-c", - HOLDER_SCRIPT, - name, - mode, - str(tmp_path / f"release-{name}"), - str(tmp_path / "events"), - ), - cwd=HARNESS_DIR, - env={**os.environ, "TMPDIR": str(tmp_path), "PYTHONPATH": str(HARNESS_DIR)}, - ) - ) - held.callback(holder.kill) - return holder - - -def test_readers_share_exclusive_waits_and_a_waiting_exclusive_beats_later_readers(tmp_path: Path) -> None: - lock_dir: Final = tmp_path / f"litellm-e2e-stack-{stack_lock.STACK_DIGEST}" - lock_dir.mkdir() - log_path: Final = tmp_path / "events" - with ExitStack() as held: - first_reader: Final = _start_holder(held, tmp_path, "A", "shared") - _wait_for_event(log_path, "A enter") - second_reader: Final = _start_holder(held, tmp_path, "R", "shared") - _wait_for_event(log_path, "R enter") - (tmp_path / "release-R").touch() - _wait_for_event(log_path, "R exit") - writer: Final = _start_holder(held, tmp_path, "W", "exclusive") - _wait_until_gate_is_held_exclusively(lock_dir / "gate") - late_reader: Final = _start_holder(held, tmp_path, "B", "shared") - _wait_for_event(log_path, "B waiting") - time.sleep(SETTLE_SECONDS) - (tmp_path / "release-A").touch() - _wait_for_event(log_path, "W enter") - (tmp_path / "release-W").touch() - _wait_for_event(log_path, "B enter") - (tmp_path / "release-B").touch() - for holder in (first_reader, second_reader, writer, late_reader): - assert holder.wait(timeout=DEADLINE_SECONDS) == 0 - events: Final = _events(log_path) - assert events.index("R enter") < events.index("A exit") - assert events.index("W enter") > events.index("A exit") - assert events.index("B enter") > events.index("W exit") diff --git a/tests/integration/pricing/test_cost_margin_discount.py b/tests/integration/pricing/test_cost_margin_discount.py index bfd2002f871..3fd1ef8694e 100644 --- a/tests/integration/pricing/test_cost_margin_discount.py +++ b/tests/integration/pricing/test_cost_margin_discount.py @@ -7,7 +7,7 @@ import httpx import pytest from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.database import read_rows -from pydantic import BaseModel, JsonValue +from pydantic import BaseModel, ConfigDict, JsonValue class CostEstimate(BaseModel): @@ -110,6 +110,25 @@ def _assert_logged_spends(expectations: tuple[tuple[str, float], ...]) -> None: assert actual == pytest.approx(dict(expectations)), rows +class LoggedBreakdown(BaseModel): + model_config = ConfigDict(frozen=True) + + original_cost: float + discount_percent: float | None = None + discount_amount: float | None = None + margin_percent: float | None = None + margin_fixed_amount: float | None = None + margin_total_amount: float | None = None + + +def _logged_breakdown(request_id: str) -> LoggedBreakdown: + rows: Final = read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)) + assert len(rows) == 1, rows + metadata: Final = rows[0]["metadata"] + parsed: Final = object_value(json.loads(metadata) if isinstance(metadata, str) else metadata) + return LoggedBreakdown.model_validate(parsed["cost_breakdown"]) + + def _estimated(gateway: Gateway, model: str) -> CostEstimate: response: Final = gateway.request( "POST", "/cost/estimate", {"model": model, "input_tokens": 20, "output_tokens": 20} @@ -236,6 +255,17 @@ def test_discount_and_margin_reprice_a_request_and_survive_rejection(gateway: Ga (streamed_global_request_id, 0.12), ) ) + discounted: Final = _logged_breakdown(discounted_request_id) + assert discounted.original_cost == pytest.approx(0.06), discounted + assert discounted.discount_percent == pytest.approx(0.05), discounted + assert discounted.discount_amount == pytest.approx(0.06 * 0.05), discounted + margined: Final = _logged_breakdown(provider_margin_request_id) + assert margined.original_cost == pytest.approx(0.06), margined + assert margined.discount_percent == pytest.approx(0.05), margined + assert margined.discount_amount == pytest.approx(0.06 * 0.05), margined + assert margined.margin_percent == pytest.approx(0.1), margined + assert margined.margin_fixed_amount == pytest.approx(0.001), margined + assert margined.margin_total_amount == pytest.approx(0.057 * 0.1 + 0.001), margined rejections: Final = ( ( diff --git a/tests/integration/spend/test_failing_request_stored_snapshot.py b/tests/integration/spend/test_failing_request_stored_snapshot.py new file mode 100644 index 00000000000..4f016c33f37 --- /dev/null +++ b/tests/integration/spend/test_failing_request_stored_snapshot.py @@ -0,0 +1,84 @@ +import json +import uuid +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment, object_value +from integration._support.process import owned_proxy +from pydantic import JsonValue + +UNREACHABLE_BASE: Final = "http://127.0.0.1:9/v1" +TRANSCRIPT_TURNS: Final = 40 +RETRIES: Final = 2 +STORED_REQUEST_BUDGET_KB: Final = 64 + + +def _config_storing_prompts(directory: Path) -> Path: + config: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + general_settings: Final = object_value(config["general_settings"]) + path: Final = directory / "store-prompts.yaml" + path.write_text( + yaml.safe_dump({**config, "general_settings": {**general_settings, "store_prompts_in_spend_logs": True}}) + ) + return path + + +@pytest.fixture(scope="module") +def gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + with gateway_from_environment() as parent: + directory: Final = tmp_path_factory.mktemp("failing_request_stored_snapshot") + with owned_proxy(parent, directory, {}, config=_config_storing_prompts(directory)) as owned: + yield owned + + +def _transcript(marker: str) -> list[JsonValue]: + turns: Final[list[JsonValue]] = [ + {"role": role, "content": f"turn {turn // 2} {role}"} + for turn, role in enumerate(("user", "assistant") * TRANSCRIPT_TURNS) + ] + return [*turns, {"role": "user", "content": marker}] + + +def _stored_request(gateway: Gateway, call_id: str) -> JsonValue: + rows: Final = eventually( + lambda: gateway.rows('SELECT proxy_server_request FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)), + lambda values: len(values) == 1, + seconds=40, + ) + stored: Final = rows[0]["proxy_server_request"] + return json.loads(stored) if isinstance(stored, str) else stored + + +def test_failing_request_with_retries_and_fallback_stores_one_copy_of_the_request(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + primary: Final = f"stored-snapshot-{uuid.uuid4().hex}" + fallback: Final = f"stored-snapshot-fb-{uuid.uuid4().hex}" + for name in (primary, primary, fallback, fallback): + scenario.model(name=name, api_base=UNREACHABLE_BASE, cooldown_time=0) + key: Final = scenario.key(models=[primary, fallback]) + marker: Final = f"stored snapshot {uuid.uuid4().hex}" + + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": primary, + "max_tokens": 16, + "messages": _transcript(marker), + "router_settings_override": {"num_retries": RETRIES, "fallbacks": [{primary: [fallback]}]}, + }, + key=key, + ) + assert response.status_code != 200, response.text + assert fallback in response.text, response.text + + serialized: Final = json.dumps(_stored_request(gateway, response.headers["x-litellm-call-id"])) + stored_kb: Final = len(serialized.encode()) / 1024 + assert marker in serialized, serialized[:400] + assert stored_kb <= STORED_REQUEST_BUDGET_KB, ( + f"one failing request stored a {stored_kb:.0f} KB request snapshot for a " + f"{2 * TRANSCRIPT_TURNS + 1}-message transcript with {RETRIES} retries and a fallback" + ) From d80ad6ebf30a8ea41a1321ebb82557e1a499d50d Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 17:29:48 -0800 Subject: [PATCH 13/41] ci(integration): keep the MCP SDK1 peer on mcp 1.x despite the project override Carries the CircleCI half of #47651 so the mcp>=2.3.0 uv override does not replace the peer venv's pinned mcp==1.28.1 (cherry picked from commit 8fe0a7358e87cad3516d005854b858b14c32f515) Co-Authored-By: Claude Opus 5.5 --- .circleci/config.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 0d442e32b0a..1edf22578ff 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1291,7 +1291,7 @@ jobs: name: Install MCP SDK1 peer command: | uv venv --python 3.12 .venv-mcp-peer - uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1' + uv pip install --no-config --exclude-newer "3 days" --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1' echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV" - when: condition: From 48bb30b5376b07a51ab346c2728fb97daa8b6a8e Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 14:29:59 -0700 Subject: [PATCH 14/41] test(mcp): replace the live Linear MCP OAuth e2e tests with integration coverage (#47154) * test(mcp): move the MCP OAuth happy path from a live Linear e2e to an integration test GitHub Actions now runs unit tests only, and the MCP OAuth happy path needed a captured Linear browser login to click through a real consent page. The OAuth flows it exercised are already integration tests against an owned authorization-server double, except one behavior: the stored per-user token must be served by a proxy process that never saw the consent, with no shared Redis to fall back on. That case is now in test_mcp_oauth_flows.py The e2e test, its owned-gateway helper, its workflow, the mcp_oauth_live marker and the JUnit gate script only that workflow used are removed * test(mcp): chat completions run an OAuth server tool with the caller's own stored token * test(e2e): drop the Linear chat-completion MCP OAuth e2e and its helpers The behavior it checked, a chat completion running an OAuth server's tool with the caller's stored per-user token, is now an integration test in test_mcp_llm_endpoints.py that needs no saved Linear browser login * test(mcp): decode the stored per-user credential as urlsafe base64 The plaintext check used b64decode, which drops the - and _ of a urlsafe value and raised on roughly one ciphertext in five * test(mcp): drop the chat OAuth integration test that duplicates test_mcp_slots F2-chat security/test_mcp_slots.py already sends a per-user OAuth token through /v1/chat/completions and asserts the peer receives it. Every mutant the new test killed was also killed by the existing suite, so it added no signal (cherry picked from commit 44fcbc30b07792aceb24e312ea5c59e5cfbdb8a9) --- .github/e2e-stack/assert_tests_ran.py | 69 --- .github/workflows/test-linting.yml | 4 +- .github/workflows/test-mcp-oauth-e2e.yml | 180 ------- .../check_e2e_no_raw_requests.py | 1 - .../test_e2e_changed_gate.py | 156 ------- .../test_e2e_junit_report.py | 16 +- tests/e2e/AGENTS.md | 8 +- tests/e2e/CONTRIBUTING.md | 52 --- tests/e2e/conftest.py | 16 - tests/e2e/coverage_registry/mcp.yaml | 8 - tests/e2e/e2e_config.py | 5 - tests/e2e/mcp/linear_session_capture.py | 58 --- tests/e2e/mcp/oauth_chat_client.py | 442 ------------------ tests/e2e/mcp/oauth_gateway.py | 197 -------- .../mcp/test_mcp_chat_completion_oauth_e2e.py | 218 --------- .../e2e/mcp/test_mcp_oauth_happy_path_e2e.py | 216 --------- tests/e2e/models.py | 38 +- tests/e2e/pytest.ini | 1 - tests/integration/mcp/test_mcp_oauth_flows.py | 48 ++ tests/unit/test_circleci_path_filter.py | 2 +- 20 files changed, 61 insertions(+), 1674 deletions(-) delete mode 100644 .github/e2e-stack/assert_tests_ran.py delete mode 100644 .github/workflows/test-mcp-oauth-e2e.yml delete mode 100644 tests/code_coverage_tests/test_e2e_changed_gate.py delete mode 100644 tests/e2e/mcp/linear_session_capture.py delete mode 100644 tests/e2e/mcp/oauth_chat_client.py delete mode 100644 tests/e2e/mcp/oauth_gateway.py delete mode 100644 tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py delete mode 100644 tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py diff --git a/.github/e2e-stack/assert_tests_ran.py b/.github/e2e-stack/assert_tests_ran.py deleted file mode 100644 index 3af49007b1e..00000000000 --- a/.github/e2e-stack/assert_tests_ran.py +++ /dev/null @@ -1,69 +0,0 @@ -import os -import re -import sys -import xml.etree.ElementTree as ET -from pathlib import Path -from typing import Final - -sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e")) -from coverage_registry.management_cases import MANAGEMENT_CASES - - -def main() -> int: - selected: Final = tuple(sys.argv[2:]) - try: - report: Final = ET.parse(Path(sys.argv[1])).getroot() - except (ET.ParseError, OSError): - _ = sys.stdout.write("::error::could not read the test execution report\n") - return 1 - cases: Final = tuple(report.iter("testcase")) - expected_count: Final = os.environ.get("E2E_REQUIRED_TEST_COUNT") - passed: Final = frozenset( - case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error")) - ) - missing: Final = tuple(path for path in selected if path not in passed) - required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected) - passed_nodes: Final = frozenset( - prop.get("value") - for case in cases - if all(case.find(tag) is None for tag in ("skipped", "failure", "error")) - for prop in case.findall("./properties/property") - if prop.get("name") == "management_node" - ) - missing_nodes: Final = required_nodes - passed_nodes - for node in sorted(missing_nodes): - _ = sys.stdout.write(f"::error::required management case did not pass: {node}\n") - for path in selected: - collected: Final = sum(case.get("file") == path for case in cases) - skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases) - _ = sys.stdout.write(f"{path}: {collected} collected, {skipped} skipped\n") - for case in cases: - if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error", "skipped")): - continue - outcome = "skipped" if case.find("skipped") is not None else "failed" - _ = sys.stdout.write(f" {outcome}: {case.get('classname', '')}::{case.get('name', '')}\n") - for prop in case.findall("./properties/property"): - name = prop.get("name", "") - value = prop.get("value", "") - if name in ("oauth_failure_phase", "oauth_exception_type", "oauth_frame") and re.fullmatch( - r"[A-Za-z0-9_.:<>-]{1,240}", value - ): - _ = sys.stdout.write(f" {name}: {value}\n") - if expected_count is not None and ( - len(cases) != int(expected_count) or any(case.find("skipped") is not None for case in cases) - ): - _ = sys.stdout.write("::error::required test count was not met or a required case was skipped\n") - return 1 - if ( - selected - and not missing - and not missing_nodes - and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error")) - ): - return 0 - _ = sys.stdout.write("::error::every selected file must execute a passing test, with no failures or errors\n") - return 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 4c8aeb64e85..3b9a2614725 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -539,8 +539,8 @@ jobs: - name: test_unit_passed_gate run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_unit_passed_gate.py - - name: test_e2e_changed_gate - run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py tests/code_coverage_tests/test_e2e_idp_stack.py + - name: test_e2e_idp_stack + run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_idp_stack.py - name: test_e2e_metadata env: diff --git a/.github/workflows/test-mcp-oauth-e2e.yml b/.github/workflows/test-mcp-oauth-e2e.yml deleted file mode 100644 index 5c07a68714e..00000000000 --- a/.github/workflows/test-mcp-oauth-e2e.yml +++ /dev/null @@ -1,180 +0,0 @@ -name: MCP OAuth happy path - -on: - pull_request: - paths: - - '.github/workflows/test-mcp-oauth-e2e.yml' - - '.github/e2e-stack/**' - - 'tests/e2e/*.py' - - 'tests/e2e/pytest.ini' - - 'tests/e2e/idp_realm.json' - - 'tests/e2e/mcp/**' - - 'litellm/experimental_mcp_client/**' - - 'litellm/proxy/_experimental/mcp_server/**' - - 'litellm/proxy/auth/**' - - 'litellm/proxy/management_endpoints/mcp_management_endpoints.py' - - 'litellm/proxy/_types.py' - - 'litellm/types/mcp_server/mcp_server_manager.py' - - 'litellm/proxy/management_endpoints/*sso*.py' - - 'litellm/proxy/management_endpoints/sso/**' - - 'litellm/proxy/common_utils/encrypt_decrypt_utils.py' - - 'litellm/proxy/proxy_server.py' - - 'litellm/proxy/schema.prisma' - - 'ui/litellm-dashboard/src/app/connect/**' - - 'ui/litellm-dashboard/src/app/mcp/oauth/**' - - 'pyproject.toml' - - 'uv.lock' - workflow_dispatch: - -permissions: {} - -concurrency: - group: mcp-oauth-${{ github.ref }} - cancel-in-progress: true - -jobs: - oauth: - if: github.event_name == 'workflow_dispatch' || github.event.pull_request.head.repo.full_name == github.repository - runs-on: ubuntu-latest - environment: e2e-changed - timeout-minutes: 45 - permissions: - contents: read - id-token: write - services: - postgres: - image: postgres:16.6 - env: - POSTGRES_USER: litellm - POSTGRES_PASSWORD: dbpassword9090 - POSTGRES_DB: litellm - ports: - - 5432:5432 - options: >- - --health-cmd "pg_isready -U litellm" - --health-interval 5s - --health-timeout 5s - --health-retries 10 - env: - DATABASE_HOST: 127.0.0.1 - DATABASE_PORT: '5432' - DATABASE_USER: litellm - DATABASE_PASSWORD: dbpassword9090 - DATABASE_NAME: litellm - DATABASE_URL: postgresql://litellm:dbpassword9090@127.0.0.1:5432/litellm - E2E_KEYCLOAK_URL: http://127.0.0.1:8081 - E2E_KEYCLOAK_ADMIN_USER: admin - E2E_KEYCLOAK_ADMIN_PASSWORD: e2e-ephemeral-idp-not-a-secret - E2E_FIXTURE_MODE: live - E2E_PROVIDER_CACHE: '0' - E2E_MCP_OAUTH_LIVE: '1' - E2E_REQUIRED_TEST_COUNT: '4' - steps: - - name: Checkout the tested source - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - ref: ${{ github.event.pull_request.head.sha || github.sha }} - persist-credentials: false - - - name: Require and materialize the upstream login - env: - STORAGE_STATE: ${{ secrets.E2E_LINEAR_STORAGE_STATE_B64 }} - run: | - umask 077 - python3 - <<'PY' - import base64 - import json - import os - import secrets - from pathlib import Path - encoded = os.environ.get("STORAGE_STATE", "") - if not encoded: - raise SystemExit("E2E_LINEAR_STORAGE_STATE_B64 is required; capture and provision a test-account login") - state = json.loads(base64.b64decode(encoded, validate=True)) - if not isinstance(state, dict) or not state.get("cookies"): - raise SystemExit("The captured login must contain browser cookies") - directory = Path(os.environ["RUNNER_TEMP"]) / "mcp-oauth-private" - directory.mkdir(mode=0o700) - path = directory / "linear-state.json" - path.write_text(json.dumps(state)) - with open(os.environ["GITHUB_ENV"], "a") as output: - output.write(f"E2E_LINEAR_STORAGE_STATE={path}\n") - for name in ("LITELLM_MASTER_KEY", "LITELLM_SALT_KEY"): - value = "sk-e2e-" + secrets.token_hex(24) - print(f"::add-mask::{value}") - output.write(f"{name}={value}\n") - PY - - - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: '3.13' - - uses: ./.github/actions/setup-uv-with-retries - with: - version: '0.10.9' - - uses: ./.github/actions/cache-cargo-build - - name: Install the frozen E2E environment - run: | - .github/scripts/uv_sync_with_retries.sh --frozen --extra proxy --extra proxy-runtime --extra extra_proxy --group ci --group proxy-dev --group e2e-dev - uv run --no-sync python scripts/prisma_generate_if_needed.py - uv run --no-sync playwright install --with-deps chromium - - - name: Configure license access - id: aws - uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0 - with: - role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }} - aws-region: us-east-1 - role-session-name: mcp-oauth-${{ github.run_id }} - role-duration-seconds: 900 - output-env-credentials: false - output-credentials: true - - name: Load the E2E license - env: - AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }} - AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }} - AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }} - AWS_DEFAULT_REGION: us-east-1 - run: | - license="$(aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license --query SecretString --output text)" - test -n "${license}" - echo "::add-mask::${license}" - echo "LITELLM_LICENSE=${license}" >> "${GITHUB_ENV}" - - - uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0 - with: - node-version-file: ui/litellm-dashboard/.nvmrc - - name: Build the gateway consent UI at the tested commit - run: | - cd ui/litellm-dashboard - ../../scripts/with_dashboard_node.sh npm ci - ../../scripts/with_dashboard_node.sh npm run build - mkdir -p ../../litellm/proxy/_experimental/out - cp -r out/. ../../litellm/proxy/_experimental/out/ - find ../../litellm/proxy/_experimental/out -name '*.html' ! -name index.html | while read -r page; do - mkdir -p "${page%.html}" - mv "${page}" "${page%.html}/index.html" - done - - - name: Prepare the isolated database and IdP - run: | - umask 077 - bash .github/e2e-stack/start-idp.sh - uv run --no-sync python migrations/run.py > "${RUNNER_TEMP}/mcp-oauth-private/migrations.log" 2>&1 - - - name: Run every required OAuth variant without retries - run: | - umask 077 - uv run --no-sync pytest -c tests/e2e/pytest.ini tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \ - --rootdir=. --reruns 0 --tb=short -o junit_family=xunit1 \ - --junitxml="${RUNNER_TEMP}/mcp-oauth-private/results.xml" \ - > "${RUNNER_TEMP}/mcp-oauth-private/pytest.log" 2>&1 - - name: Report JUnit results and reject skipped or missing cases - if: always() - run: | - uv run --no-sync python .github/e2e-stack/assert_tests_ran.py \ - "${RUNNER_TEMP}/mcp-oauth-private/results.xml" tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py - - name: Remove private login and logs - if: always() - run: | - docker rm -f e2e-keycloak >/dev/null 2>&1 || true - rm -rf "${RUNNER_TEMP}/mcp-oauth-private" diff --git a/tests/code_coverage_tests/check_e2e_no_raw_requests.py b/tests/code_coverage_tests/check_e2e_no_raw_requests.py index 4316f1b57b6..bea177d3b7d 100644 --- a/tests/code_coverage_tests/check_e2e_no_raw_requests.py +++ b/tests/code_coverage_tests/check_e2e_no_raw_requests.py @@ -25,7 +25,6 @@ ALLOWED_RAW_CLIENT_FILES = { "e2e/e2e_http.py": ("requests",), "e2e/conftest.py": ("requests",), "e2e/claude_code/pr_gate_version_resolver.py": ("urllib.request",), - "e2e/mcp/oauth_chat_client.py": ("httpx",), } EXCEPTION_ONLY_NAMES = frozenset({"RequestException", "ConnectionError", "Timeout", "HTTPError"}) diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py deleted file mode 100644 index c4d12d1282c..00000000000 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ /dev/null @@ -1,156 +0,0 @@ -import os -import subprocess -import sys -import xml.etree.ElementTree as ET -from pathlib import Path -from typing import Final - -import pytest - -GATE: Final = Path(__file__).resolve().parents[2] / ".github/e2e-stack/assert_tests_ran.py" -SELECTED: Final = ("tests/e2e/access_control/test_a.py", "tests/e2e/access_control/test_b.py") - - -@pytest.mark.parametrize( - ("second_outcome", "expected_status"), - (("passed", 0), ("skipped", 1), ("failure", 1), ("error", 1), ("deselected", 1)), -) -def test_each_changed_file_must_run(tmp_path: Path, second_outcome: str, expected_status: int) -> None: - suite: Final = ET.Element("testsuite") - _ = ET.SubElement(suite, "testcase", file=SELECTED[0]) - if second_outcome != "deselected": - second: Final = ET.SubElement(suite, "testcase", file=SELECTED[1]) - if second_outcome != "passed": - _ = ET.SubElement(second, second_outcome) - report: Final = tmp_path / "report.xml" - ET.ElementTree(suite).write(report) - - result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True) - - assert result.returncode == expected_status, result.stdout - - -@pytest.mark.parametrize("outcome", ("failure", "error")) -def test_passing_case_does_not_hide_a_failure_in_the_same_file(tmp_path: Path, outcome: str) -> None: - suite: Final = ET.Element("testsuite") - _ = ET.SubElement(suite, "testcase", file=SELECTED[0]) - failed: Final = ET.SubElement(suite, "testcase", file=SELECTED[0]) - _ = ET.SubElement(failed, outcome) - report: Final = tmp_path / "report.xml" - ET.ElementTree(suite).write(report) - - result: Final = subprocess.run( - [sys.executable, str(GATE), str(report), SELECTED[0]], capture_output=True, text=True - ) - - assert result.returncode == 1 - - -def test_failed_cases_are_named_per_selected_file(tmp_path: Path) -> None: - suite: Final = ET.Element("testsuite") - _ = ET.SubElement(suite, "testcase", file=SELECTED[0], classname="tests.e2e.access_control.test_a", name="test_ok") - failed: Final = ET.SubElement( - suite, "testcase", file=SELECTED[0], classname="tests.e2e.access_control.test_a", name="test_boom" - ) - _ = ET.SubElement(failed, "failure", message="secret-bearing message") - errored: Final = ET.SubElement( - suite, "testcase", file=SELECTED[1], classname="tests.e2e.access_control.test_b", name="test_setup" - ) - _ = ET.SubElement(errored, "error") - report: Final = tmp_path / "report.xml" - ET.ElementTree(suite).write(report) - - result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True) - - assert result.returncode == 1 - assert " failed: tests.e2e.access_control.test_a::test_boom\n" in result.stdout - assert " failed: tests.e2e.access_control.test_b::test_setup\n" in result.stdout - assert "test_ok" not in result.stdout - assert "secret-bearing message" not in result.stdout - - -@pytest.mark.parametrize("contents", ("", "')) -def test_missing_execution_evidence_fails(tmp_path: Path, contents: str) -> None: - report: Final = tmp_path / "report.xml" - _ = report.write_text(contents) - - result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True) - - assert result.returncode == 1 - - -@pytest.mark.parametrize("omitted_role", ("proxy_admin", "team_member", "internal_user_viewer")) -def test_one_passing_management_case_cannot_hide_a_missing_actor(tmp_path: Path, omitted_role: str) -> None: - suite: Final = ET.Element("testsuite") - path: Final = "tests/e2e/management/test_jwt_management_e2e.py" - case: Final = ET.SubElement(suite, "testcase", file=path) - properties: Final = ET.SubElement(case, "properties") - _ = ET.SubElement( - properties, - "property", - name="management_node", - value=f"{path}::TestJwtManagement::test_actor_subject_and_database_role[proxy_admin_viewer]", - ) - report: Final = tmp_path / "report.xml" - ET.ElementTree(suite).write(report) - result: Final = subprocess.run([sys.executable, str(GATE), str(report), path], capture_output=True, text=True) - assert result.returncode == 1 - assert f"test_actor_subject_and_database_role[{omitted_role}]" in result.stdout - - -@pytest.mark.parametrize("phase", ("setup", "call", "teardown")) -@pytest.mark.parametrize("required_count", ("1", "4")) -def test_oauth_failure_diagnostics_do_not_publish_private_payloads( - tmp_path: Path, phase: str, required_count: str -) -> None: - suite: Final = ET.Element("testsuite") - case: Final = ET.SubElement(suite, "testcase", file=SELECTED[0]) - private: Final = "private-token-in-exception-message" - failure: Final = ET.SubElement(case, "failure", message=private) - failure.text = private - properties: Final = ET.SubElement(case, "properties") - for name, value in ( - ("oauth_failure_phase", phase), - ("oauth_exception_type", "AssertionError"), - ("oauth_frame", "oauth_gateway.py:120:start"), - ("oauth_frame", f"injected\\n{private}"), - ("unrelated_property", private), - ): - _ = ET.SubElement(properties, "property", name=name, value=value) - report: Final = tmp_path / "report.xml" - ET.ElementTree(suite).write(report) - result: Final = subprocess.run( - [sys.executable, "-I", str(GATE), str(report), SELECTED[0]], - capture_output=True, - text=True, - env={**os.environ, "E2E_REQUIRED_TEST_COUNT": required_count}, - ) - assert result.returncode == 1 - assert f"oauth_failure_phase: {phase}" in result.stdout - assert "oauth_exception_type: AssertionError" in result.stdout - assert "oauth_frame: oauth_gateway.py:120:start" in result.stdout - assert private not in result.stdout + result.stderr - - -@pytest.mark.parametrize( - ("count", "skip", "expected"), ((0, False, 1), (3, False, 1), (4, False, 0), (5, False, 1), (4, True, 1)) -) -def test_required_count_reports_cases_before_rejecting(tmp_path: Path, count: int, skip: bool, expected: int) -> None: - suite = ET.Element("testsuite") - for index in range(count): - case = ET.SubElement(suite, "testcase", file=SELECTED[0], classname="OAuth", name=f"variant{index}") - if skip and index == 0: - ET.SubElement(case, "skipped", message="private-skip-reason") - report = tmp_path / "report.xml" - ET.ElementTree(suite).write(report) - result = subprocess.run( - [sys.executable, "-I", str(GATE), str(report), SELECTED[0]], - env={**os.environ, "E2E_REQUIRED_TEST_COUNT": "4"}, - capture_output=True, - text=True, - ) - assert result.returncode == expected - assert f"{count} collected, {int(skip)} skipped" in result.stdout - if skip: - assert "skipped: OAuth::variant0" in result.stdout - assert "private-skip-reason" not in result.stdout + result.stderr diff --git a/tests/code_coverage_tests/test_e2e_junit_report.py b/tests/code_coverage_tests/test_e2e_junit_report.py index 140b9ba6dca..be5fcd245db 100644 --- a/tests/code_coverage_tests/test_e2e_junit_report.py +++ b/tests/code_coverage_tests/test_e2e_junit_report.py @@ -172,12 +172,7 @@ def open_consent() -> None: raise RuntimeError("consent page timed out") -@pytest.mark.mcp_oauth_live -def test_oauth_dies_on_consent() -> None: - open_consent() - - -def test_plain_dies_on_consent() -> None: +def test_dies_on_consent() -> None: open_consent() """ @@ -353,16 +348,13 @@ class TestStepsReachTheReport: def test_a_failed_phase_s_own_report_carries_the_steps(tmp_path: Path) -> None: """Plugins that read the failed setup or call report, not the teardown one - junitxml writes from, see where the test died too, oauth-live or not.""" + junitxml writes from, see where the test died too.""" write_suite(tmp_path, {"test_consent.py": FAILED_PHASE_SUITE, "report_spy.py": REPORT_SPY_PLUGIN}) - child: Final = run_child_pytest(tmp_path, "-p", "report_spy", env={"E2E_MCP_OAUTH_LIVE": "1"}) + child: Final = run_child_pytest(tmp_path, "-p", "report_spy") seen_path: Final = tmp_path / "failed-reports.jsonl" assert seen_path.exists(), f"no failed report reached the spy:\n{child.stdout}\n{child.stderr}" seen: Final = dict(map(FailedReport.validate_json, seen_path.read_text().splitlines())) - assert seen == { - "test_oauth_dies_on_consent": ("open the consent page",), - "test_plain_dies_on_consent": ("open the consent page",), - }, child.stdout + assert seen == {"test_dies_on_consent": ("open the consent page",)}, child.stdout class TestDeclaredPropertiesReachTheReport: diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 53e203a4d7b..88a4df4c40c 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -40,7 +40,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) - `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface, not Playwright) - `a2a/` - the A2A (agent-to-agent) surface: admin registration via `/v1/agents`, proxy-fronted card discovery at `/.well-known/agent-card.json`, and JSON-RPC `message/send` invocation, driving agents backed by the litellm completion bridge (a real provider) and asserting protocol-version normalization (0.3 vs 1.0) -- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below), plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions` in `test_mcp_chat_completion_oauth_e2e.py` and direct MCP protocol operations in `test_mcp_oauth_happy_path_e2e.py`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion or protocol call lists and executes the server's tools with the stored per-user token. The semantic tool filter test turns on a proxy-wide setting, so it is marked `proxy_wide_config` and runs on the redeployed managed-files stack after the main run, never next to it +- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below). The semantic tool filter test turns on a proxy-wide setting, so it is marked `proxy_wide_config` and runs on the redeployed managed-files stack after the main run, never next to it - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) @@ -53,14 +53,14 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family ## MCP suite: real Datadog only -Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datadog remote MCP server, except the two Linear OAuth tests `test_mcp_chat_completion_oauth_e2e.py` and `test_mcp_oauth_happy_path_e2e.py`. Do not add a compose service, FastMCP fixture, mock upstream, or any other fake MCP host for this suite +Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datadog remote MCP server. Do not add a compose service, FastMCP fixture, mock upstream, or any other fake MCP host for this suite - Register via `register_datadog_mcp` in `tests/e2e/mcp/datadog_mcp.py` (or extend that helper if you need a different `toolsets=` / `allowed_tools` slice of the same Datadog endpoint). That posts `/v1/mcp/server` with `url=datadog_mcp_url(...)` and static headers `DD-API-KEY` / `DD-APPLICATION-KEY` from the process env - Auth is Datadog's documented CI/header path, not a browser OAuth authorize/token dance. Hard-fail when `DD_API_KEY` or `DD_APP_KEY` is missing (`assert_dd_mcp_creds`). Never skip for a missing fake upstream - Prefer calling real Datadog tools that prove the product path (e.g. `search_datadog_logs` for list/call and permission denials). Seed a unique marker (`e2e-datadog-mcp-*`) in a chat completion when you need a log the tool can find. Dual-read with `dd_logs` from conftest when delivery matters - Delete the MCP server (and any keys) through `resources.defer` the same way every other suite tears down - If a new MCP behavior cannot be covered with Datadog's tool surface, say so in the PR and get agreement before inventing another upstream. The default is always Datadog -- The two standing exceptions are `test_mcp_chat_completion_oauth_e2e.py` and `test_mcp_oauth_happy_path_e2e.py`. Datadog authenticates with the static `DD-API-KEY` / `DD-APPLICATION-KEY` headers and exposes no authorize/token dance at all, so these tests drive a real Linear MCP server instead. They are still real remote upstreams, so the no-mock, no-fixture rule above holds unchanged. The direct OAuth test also uses the existing live provider edge to inspect forwarded headers without replay, and owns a separate source-built gateway for cold restarts +- Gateway-managed OAuth (an authorize/token dance and per-user stored tokens) is out of scope here, since Datadog only takes static headers. Those flows are integration tests against an owned authorization-server double, in `tests/integration/mcp/test_mcp_oauth_flows.py` and, for a chat completion calling the tool with the caller's stored token, the F2-chat cases of `tests/integration/security/test_mcp_slots.py`. Do not add a third-party OAuth upstream that needs a saved browser login ## Lay the pattern down in a class @@ -243,7 +243,7 @@ MCPs - endpoint features with the protocol op as the variant mcp... operation : list_tools | call_tool | list_resources | read_resource | list_prompts | get_prompt auth_family : none | api_key | bearer | oauth - assertion : succeeds | denied_without_permission | persists_across_processes + assertion : succeeds | denied_without_permission e.g. mcp.call_tool.oauth.succeeds ``` diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index f85a6b58b86..4f0c7624e31 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -250,55 +250,3 @@ The semantic header set is `content-type`, `accept`, `anthropic-version`, `anthr Excluded transport and telemetry headers are `host`, `content-length`, `connection`, `accept-encoding`, `user-agent`, `traceparent`, `tracestate`, `x-request-id`, `x-client-request-id` and `x-stainless-*`. Inbound transfer-encoding is unsupported; send JSON with content-length framing. The destination represents host identity and the relay carries original body bytes. Replay does not verify credentials, SDK timeout/retry behavior, transport performance, model availability or stateful remote IDs. Live relay uses original request bytes and header values, never the stored identity Strict replay harness regression tests live in `tests/code_coverage_tests/test_provider_replay_harness.py`. The `python` job in `.github/workflows/test-linting.yml` runs them alongside the provider-edge and fixture tests at the root of `tests/e2e_harness/` with `--noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures`; they need only synthetic HTTP providers and temporary fixture storage - - -## MCP OAuth happy path - -`test_mcp_oauth_happy_path_e2e.py` runs one shared scenario with four variants: -aggregate gateway SSO and explicitly configured per-server JWT, each directly -against Linear and through the live provider edge. The edge forwards to real -Linear without replay and compares the forwarded bearer to the encrypted -canonical user/server credential. This observes the forwarding boundary, not -Linear's internal logs. Direct variants independently exercise discovery - -Use the existing database preparation, Prisma generation and Keycloak setup. -Build and stage the dashboard from the tested checkout as in the UI runner. -Provide `DATABASE_URL`, `LITELLM_MASTER_KEY`, `LITELLM_SALT_KEY`, `LITELLM_LICENSE`, -and the `E2E_KEYCLOAK_*` settings. Capture a test-account Linear login using -`mcp/linear_session_capture.py` and set `E2E_LINEAR_STORAGE_STATE` to that private -file. The test workspace must contain a team. Do not publish browser state or -raw test/proxy output - -```bash -E2E_MCP_OAUTH_LIVE=1 E2E_FIXTURE_MODE=live E2E_PROVIDER_CACHE=0 \ - uv run --no-sync pytest tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \ - --rootdir=. --reruns 0 -``` - -The test starts and restarts its own source-built proxy on a free loopback port, -retaining its database and SSO client but no Redis or process-local cache. It -does not restart an existing proxy or clear shared databases. Gateway login, -consent, immediate list/call and post-restart reconnect must all succeed. The -aggregate client never injects a gateway header; the explicitly labeled JWT -variant configures `x-litellm-api-key` for the first consent and reconnects with -only its gateway JWT after restart - -`.github/workflows/test-mcp-oauth-e2e.yml` automatically requests a run for -same-repository pull requests changing MCP, gateway authentication/SSO, consent -UI, dependencies or the relevant E2E harness/workflow paths. It retains manual -`workflow_dispatch` for targeted verification. The four cases run in the -protected `e2e-changed` environment after its normal deployment approval; -reviewers should approve and inspect this separate OAuth check when it appears. -Fork pull requests do not run this credentialed job; use a reviewed -same-repository branch for their verification. The workflow's path-filtered -check is not configured here as a globally required branch-protection check. -Provision `E2E_LINEAR_STORAGE_STATE_B64` as a secret there and retain the existing E2E license/AWS role configuration. A missing or -expired session fails the job; collection, deselection and skips are not passes. -No LLM call is needed - -Coverage remains limited to authorization-code OAuth over HTTP. M2M, OBO, -PKCE passthrough, static/BYOK, ID-JAG, forwarding, SigV4 and stdio are outside this -scenario; consult the registry and LIT-3559 for their existing coverage and gaps. -LIT-4506 owns broader isolation/failure regressions. LIT-7737 retains ownership -of dependency/Python compatibility and its matrix; this test reuses its delivered -environment and does not change dependency constraints or compatibility gates diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 309168351e7..8fb96bef2f9 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -17,7 +17,6 @@ import functools import os from collections.abc import Generator, Iterator from datetime import datetime, timezone -from pathlib import Path from types import MappingProxyType from typing import Final @@ -29,7 +28,6 @@ from e2e_config import ( FIXTURE_DIR, FIXTURE_MODE_RAW, MANAGED_FILES_OPT_IN_ENV, - MCP_OAUTH_LIVE_OPT_IN_ENV, OTEL_TLS_OPT_IN_ENV, OTEL_V2_OPT_IN_ENV, OWNED_GATEWAY_OPT_IN_ENV, @@ -72,7 +70,6 @@ OPT_IN_MARKERS: Final = MappingProxyType( "redis_chaos": PERF_TESTS_OPT_IN_ENV, "perf": PERF_TESTS_OPT_IN_ENV, "cli_determinism": CLI_DETERMINISM_OPT_IN_ENV, - "mcp_oauth_live": MCP_OAUTH_LIVE_OPT_IN_ENV, "provider_edge_host": PROVIDER_EDGE_HOST_OPT_IN_ENV, "owned_gateway": OWNED_GATEWAY_OPT_IN_ENV, "otel_v2": OTEL_V2_OPT_IN_ENV, @@ -166,11 +163,6 @@ def pytest_configure(config: pytest.Config) -> None: "redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from " "gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set", ) - config.addinivalue_line( - "markers", - "mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless " - "E2E_MCP_OAUTH_LIVE is set", - ) config.addinivalue_line( "markers", "provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the " @@ -326,8 +318,6 @@ def pytest_runtest_setup(item: pytest.Item) -> None: item.user_properties.extend(item.config.stash[_IDLE_RSS].junit_properties) if not _reaches_proxy(item): return - if isinstance(item, pytest.Function) and "oauth_gateway" in item.fixturenames: - return reason = _proxy_fail_reason() if reason is not None: pytest.fail(reason) @@ -371,12 +361,6 @@ def pytest_runtest_makereport( report = yield if report.when in ("setup", "call"): attach_step_properties(item) - if item.get_closest_marker("mcp_oauth_live") is not None and call.excinfo is not None: - # Publish code locations only, never exception messages, source text or locals. - item.user_properties.append(("oauth_failure_phase", report.when)) - item.user_properties.append(("oauth_exception_type", call.excinfo.type.__name__)) - for entry in call.excinfo.traceback: - item.user_properties.append(("oauth_frame", f"{Path(entry.path).name}:{entry.lineno + 1}:{entry.name}")) if report.when == "call": item.stash[_CALL_PASSED] = report.passed report.user_properties = list(item.user_properties) diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index cdd194b267f..2ad815930a9 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -79,14 +79,6 @@ assertions: [succeeds] source: "db.py user_oauth_credential lookup" rationale: OAuth2 token passthrough; per-user credential storage -- id: mcp.call_tool.oauth.persists_across_processes - module: mcp - tier: P1 - operation: call_tool - auth_family: oauth - assertions: [persists_across_processes] - source: "outbound_credentials/per_user_oauth_store.py V2PerUserTokenStore" - rationale: Stored per-user token survives a verified restart of an owned gateway with no Redis cache - id: mcp.list_tools.none.succeeds module: mcp tier: P1 diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index bcab7e35ae7..6316d22d3fa 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -111,10 +111,6 @@ CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haik CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5") S3_PARTITION_GRANULARITY = os.environ.get("E2E_S3_PARTITION_GRANULARITY", "day") -LINEAR_MCP_URL = os.environ.get("E2E_LINEAR_MCP_URL", "https://mcp.linear.app/mcp") -LINEAR_STORAGE_STATE = os.environ.get("E2E_LINEAR_STORAGE_STATE", "") -LINEAR_READONLY_TOOL: Final = "list_teams" # as listed by tools/list on mcp.linear.app when PR #33787 landed - # Jaeger query API of the compose stack's OTEL trace destination (the `jaeger` # service in docker-compose.yml maps it to host 16686). Trace-completeness tests # read exported spans back through it. @@ -207,7 +203,6 @@ MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" PERF_TESTS_OPT_IN_ENV: Final = "E2E_PERF_TESTS" CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" -MCP_OAUTH_LIVE_OPT_IN_ENV: Final = "E2E_MCP_OAUTH_LIVE" PROVIDER_EDGE_HOST_OPT_IN_ENV: Final = "E2E_PROVIDER_EDGE_HOST_REACHABLE" OWNED_GATEWAY_OPT_IN_ENV: Final = "E2E_OWNED_GATEWAY" OTEL_V2_OPT_IN_ENV: Final = "E2E_OTEL_V2" diff --git a/tests/e2e/mcp/linear_session_capture.py b/tests/e2e/mcp/linear_session_capture.py deleted file mode 100644 index 1c867e17e22..00000000000 --- a/tests/e2e/mcp/linear_session_capture.py +++ /dev/null @@ -1,58 +0,0 @@ -"""One-time helper to capture a logged-in Linear browser session for the -real-Linear MCP e2e test. - -The real-Linear test drives the genuine gateway-managed authorization_code -dance against ``mcp.linear.app``. The only step that cannot be scripted is -Linear's login (magic link / SSO), so a human authenticates once here and the -resulting session (cookies + local storage) is persisted to disk. The e2e test -then loads that session in a headless Playwright context and clicks Approve on -Linear's consent screen every run, with no human and no login automation. - -Run it with the e2e venv, log into Linear in the window that opens, then return -to the terminal and press Enter: - - LITELLM=~/litellm-mcpe2e - "$LITELLM"/.venv/bin/python "$LITELLM"/tests/e2e/mcp/linear_session_capture.py - -The session is written to ``E2E_LINEAR_STORAGE_STATE`` (default -``~/.litellm-e2e/linear_storage_state.json``), outside the repo. It is a -secret: never commit it. Re-run this whenever Linear expires the session. -""" - -from __future__ import annotations - -import os -from pathlib import Path - -from playwright.sync_api import sync_playwright - -DEFAULT_STATE_PATH = Path.home() / ".litellm-e2e" / "linear_storage_state.json" - - -def capture(state_path: Path) -> None: - """Open a headed browser at Linear, wait for the human to log in, then save - the authenticated session to ``state_path``.""" - state_path.parent.mkdir(parents=True, exist_ok=True) - with sync_playwright() as playwright: - browser = playwright.chromium.launch(headless=False) - context = browser.new_context() - page = context.new_page() - page.goto("https://linear.app/login", wait_until="domcontentloaded") - print("\n" + "=" * 72) - print("Log into Linear in the browser window that just opened.") - print("If Linear emails you a magic link, paste the link into THIS window's") - print("address bar (opening it in your default browser won't capture the") - print("session). Google SSO works too as long as you complete it here.") - print("When your Linear workspace has loaded, come back and press Enter.") - print("=" * 72) - input("Press Enter once you are logged in... ") - page.goto("https://mcp.linear.app/", wait_until="domcontentloaded") - context.storage_state(path=str(state_path)) - browser.close() - print(f"\nSaved Linear session to {state_path}") - print("Point the e2e test at it with:") - print(f' export E2E_LINEAR_STORAGE_STATE="{state_path}"') - - -if __name__ == "__main__": - capture(Path(os.environ.get("E2E_LINEAR_STORAGE_STATE", str(DEFAULT_STATE_PATH)))) diff --git a/tests/e2e/mcp/oauth_chat_client.py b/tests/e2e/mcp/oauth_chat_client.py deleted file mode 100644 index 68e838d8745..00000000000 --- a/tests/e2e/mcp/oauth_chat_client.py +++ /dev/null @@ -1,442 +0,0 @@ -"""Client for the mcp chat-completion OAuth e2e suite. - -Registers a gateway-managed OAuth (authorization_code) MCP server, seeds the -per-user upstream token by driving the interactive authorize dance with the -official mcp SDK's OAuthClientProvider (the browser leg is a headless Chromium -primed with a human's saved Linear session), then exercises the server through -/chat/completions, where the gateway lists and executes its tools with the -stored per-user token. - -Management routes (/v1/mcp/server CRUD, /chat/completions) go through the -shared ProxyClient transport. The MCP protocol used to seed the token goes through -the mcp SDK, the same library production MCP hosts run. -""" - -from __future__ import annotations - -import asyncio -import re -import time -from dataclasses import dataclass -from typing import TYPE_CHECKING, Final -from urllib.parse import parse_qsl - -import httpx -import httpx2 -import pytest -from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT -from e2e_http import AuthHeaders, NoBody, unwrap -from e2e_metadata import step -from idp import Identity -from mcp import ClientSession -from mcp.client.auth import OAuthClientProvider -from mcp.client.streamable_http import streamable_http_client -from mcp.shared.auth import AuthorizationCodeResult, OAuthClientInformationFull, OAuthClientMetadata, OAuthToken -from mcp.types import TextContent -from models import ( - ChatBody, - ChatResponse, - McpOauthUserCredentialStatus, - McpServerCreateBody, - McpServerInfo, - McpServerUserCredentialListResponse, - McpServerUserCredentialRow, -) -from proxy_client import ProxyClient - -if TYPE_CHECKING: - from playwright.async_api import Route - -# Where the "browser" lands at the end of the authorize dance. Nothing listens -# here: the route interceptor short-circuits the final redirect and reads the -# code/state off its query string, exactly like a desktop MCP host intercepting -# its loopback redirect. -OAUTH_CLIENT_REDIRECT_URI = "http://127.0.0.1:53682/e2e/callback" -BROWSER_CONSENT_TIMEOUT = 60.0 - - -def _mcp_url(alias: str, base_url: str = PROXY_BASE_URL) -> str: - return f"{base_url.rstrip('/')}/{alias}/mcp" - - -class InMemoryTokenStorage: - """The mcp SDK's TokenStorage protocol, in memory for one dance: the - DCR-registered client and the gateway tokens minted for it.""" - - def __init__(self) -> None: - self._tokens: OAuthToken | None = None - self._client_info: OAuthClientInformationFull | None = None - - async def get_tokens(self) -> OAuthToken | None: - return self._tokens - - async def set_tokens(self, tokens: OAuthToken) -> None: - self._tokens = tokens - - async def get_client_info(self) -> OAuthClientInformationFull | None: - return self._client_info - - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: - self._client_info = client_info - - -async def _browser_follow_authorize( - start_url: str, - storage_state_path: str, - identity: Identity | None = None, - server_alias: str | None = None, - allow_upstream_consent: bool = True, -) -> tuple[str, str | None]: - """Play the browser's role for a real upstream whose authorize endpoint - serves an interactive consent page (Linear). A headless Chromium primed - with a human's saved Linear session opens the gateway authorize URL and - clicks through Linear's consent screens (the mcp.linear.app Approve form, - then the linear.app workspace-selection page), riding the rest of the chain - (Linear -> gateway callback -> host redirect_uri). The final hop is - intercepted and short-circuited, since nothing listens there, and its - code/state are read off the query string.""" - from playwright.async_api import async_playwright - - captured: dict[str, str] = {} # mutable-ok: hand-off from the request listener - trail: list[str] = [] # mutable-ok: navigation diagnostics for a failed dance - - def _note_request(request: object) -> None: - url = getattr(request, "url", "") - host = httpx.URL(url).host - if not allow_upstream_consent and (host == "linear.app" or host.endswith(".linear.app")): - captured["upstream_consent"] = "seen" - if url.startswith(OAUTH_CLIENT_REDIRECT_URI) and "url" not in captured: - captured["url"] = url - - async def _swallow_redirect(route: Route) -> None: - await route.fulfill(status=200, content_type="text/plain", body="ok") - - async with async_playwright() as playwright: - browser = await playwright.chromium.launch(headless=True) - context = await browser.new_context(storage_state=storage_state_path) - await context.route(re.compile(re.escape(OAUTH_CLIENT_REDIRECT_URI) + r".*"), _swallow_redirect) - page = await context.new_page() - context.on("request", _note_request) - page.on("framenavigated", lambda frame: trail.append(frame.url.split("?", 1)[0])) - await page.goto(start_url, wait_until="domcontentloaded") - deadline = time.monotonic() + BROWSER_CONSENT_TIMEOUT - while "url" not in captured and time.monotonic() < deadline: - try: - await page.wait_for_load_state("networkidle", timeout=8000) - except Exception: # noqa: BLE001 - a busy consent page never idles; fall through and try to advance it - pass - if "upstream_consent" in captured or "url" in captured: - break - if await page.locator("#username").count() and identity is not None: - await page.locator("#username").fill(identity.username) - await page.locator("#password").fill(identity.password) - await page.locator("#kc-login").click() - continue - if "/ui/connect" in page.url and server_alias is not None: - card = page.locator("div.cursor-pointer").filter(has=page.get_by_text(server_alias, exact=True)) - if await card.count() != 1: - await asyncio.sleep(0.5) - continue - connect = card.get_by_text("Connect", exact=True) - if await connect.count(): - await connect.click() - continue - if not await card.locator("svg.text-success").count(): - await asyncio.sleep(0.5) - continue - finish = page.get_by_role("button", name="Finish connecting", exact=True) - if await finish.count() and await finish.is_enabled(): - await finish.click() - continue - control = page.locator( - 'button[name="action"][value="approve"], button:has-text("Authorize"), ' - 'button:has-text("Allow"), button:has-text("@"), a:has-text("@")' - ).first - try: - await control.click(timeout=5000) - except Exception: # noqa: BLE001 - nothing to advance yet; loop and re-check - await asyncio.sleep(0.5) - final_url = page.url - await browser.close() - - # A redirect chain can finish inside goto/networkidle before the loop checks the page. - assert "upstream_consent" not in captured, "cold reconnect required upstream consent" - landing = captured.get("url") - assert landing is not None, ( - f"consent flow never reached {OAUTH_CLIENT_REDIRECT_URI}; " - f"final={final_url.split('?', 1)[0]!r}; trail={trail[-6:]}" - ) - params = dict(parse_qsl(httpx.URL(landing).query.decode())) - assert "code" in params, "client redirect_uri carried no authorization code" - return params["code"], params.get("state") - - -def _oauth_provider( - url: str, - storage: InMemoryTokenStorage, - storage_state_path: str | None, - identity: Identity | None = None, - server_alias: str | None = None, - allow_upstream_consent: bool = True, -) -> OAuthClientProvider: - """The SDK's real OAuth machinery (RFC 9728/8414 discovery, RFC 7591 DCR, - PKCE, token exchange) with the browser leg driven by Playwright against the - upstream's consent screen.""" - code_holder: dict[str, str | None] = {} # mutable-ok: hand-off between the two SDK callbacks - - async def _reject_redirect(_: str) -> None: - raise AssertionError("gateway demanded a fresh upstream consent; stored per-user token was not reused") - - async def _follow_redirect(authorize_url: str) -> None: - assert storage_state_path is not None - code, state = await _browser_follow_authorize( - authorize_url, storage_state_path, identity, server_alias, allow_upstream_consent - ) - code_holder["code"] = code - code_holder["state"] = state - - redirect_handler: Final = _reject_redirect if storage_state_path is None else _follow_redirect - - async def callback_handler() -> AuthorizationCodeResult: - code = code_holder.get("code") - assert code is not None, "callback_handler ran before the authorize redirect completed" - return AuthorizationCodeResult(code=code, state=code_holder.get("state")) - - return OAuthClientProvider( - server_url=url, - client_metadata=OAuthClientMetadata.model_validate( - { - "redirect_uris": [OAUTH_CLIENT_REDIRECT_URI], - "token_endpoint_auth_method": "none", - "grant_types": ["authorization_code", "refresh_token"], - "response_types": ["code"], - "client_name": "e2e-mcp-host", - } - ), - storage=storage, - redirect_handler=redirect_handler, - callback_handler=callback_handler, - ) - - -class _HeaderInjectingTransport(httpx2.AsyncBaseTransport): - """Adds the caller's LiteLLM key header to every outgoing SDK request - (discovery, DCR, token exchange), so the gateway resolves which user to - store the upstream token for from the key on the token exchange, exactly - like a production MCP host configured with a LiteLLM key header.""" - - def __init__(self, inner: httpx2.AsyncBaseTransport, headers: dict[str, str], gateway_url: str) -> None: - self._inner = inner - self._headers = headers - self._gateway_url = httpx2.URL(gateway_url) - - @staticmethod - def _port(url: httpx2.URL) -> int | None: - if url.port is not None: - return url.port - return {"http": 80, "https": 443}.get(url.scheme) - - async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response: - same_origin: Final = ( - request.url.scheme == self._gateway_url.scheme - and request.url.host == self._gateway_url.host - and self._port(request.url) == self._port(self._gateway_url) - ) - if same_origin: - for name, value in self._headers.items(): - if name not in request.headers: - request.headers[name] = value - else: - for name, value in self._headers.items(): - if request.headers.get(name) == value: - del request.headers[name] - return await self._inner.handle_async_request(request) - - async def aclose(self) -> None: - await self._inner.aclose() - - -def _oauth_http_client( - headers: dict[str, str], auth: OAuthClientProvider, gateway_url: str = PROXY_BASE_URL -) -> httpx2.AsyncClient: - return httpx2.AsyncClient( - auth=auth, - timeout=httpx2.Timeout(REQUEST_TIMEOUT), - follow_redirects=True, - transport=_HeaderInjectingTransport(httpx2.AsyncHTTPTransport(), headers, gateway_url), - ) - - -async def _seed_via_dance( - url: str, headers: dict[str, str], storage: InMemoryTokenStorage, storage_state_path: str -) -> tuple[str, ...]: - async with _oauth_http_client(headers, _oauth_provider(url, storage, storage_state_path)) as http_client: - async with streamable_http_client(url, http_client=http_client) as (read, write): - async with ClientSession(read, write) as session: - await session.initialize() - listed = await session.list_tools() - return tuple(sorted(tool.name for tool in listed.tools)) - - -@dataclass(frozen=True, slots=True) -class OauthToolRun: - tools: tuple[str, ...] - is_error: bool - text: str - - -async def _list_and_call( - url: str, - headers: dict[str, str], - storage: InMemoryTokenStorage, - storage_state_path: str | None, - tool: str, - arguments: dict[str, str], - gateway_url: str = PROXY_BASE_URL, - identity: Identity | None = None, - server_alias: str | None = None, - allow_upstream_consent: bool = True, -) -> OauthToolRun: - async with _oauth_http_client( - headers, - _oauth_provider(url, storage, storage_state_path, identity, server_alias, allow_upstream_consent), - gateway_url, - ) as http_client: - async with streamable_http_client(url, http_client=http_client) as (read, write): - async with ClientSession(read, write) as session: - await session.initialize() - listed: Final = await session.list_tools() - result: Final = await session.call_tool(tool, arguments) - text: Final = "".join(content.text for content in result.content if isinstance(content, TextContent)) - return OauthToolRun( - tools=tuple(sorted(tool_item.name for tool_item in listed.tools)), - is_error=result.is_error, - text=text, - ) - - -@dataclass(frozen=True, slots=True) -class ChatMcpClient: - proxy: ProxyClient - - @step("Register the MCP server with the alias {body.alias}") - def create_server(self, body: McpServerCreateBody) -> McpServerInfo: - return unwrap( - self.proxy.transport.post( - "/v1/mcp/server", - headers=self.proxy.transport.master, - json=body, - response_type=McpServerInfo, - ) - ) - - @step("Read the MCP server back from /v1/mcp/server") - def server_info(self, server_id: str) -> McpServerInfo: - return unwrap( - self.proxy.transport.get( - f"/v1/mcp/server/{server_id}", - headers=self.proxy.transport.master, - params=NoBody(), - response_type=McpServerInfo, - ) - ) - - @step("Delete the MCP server") - def delete_server(self, server_id: str) -> None: - _ = self.proxy.transport.delete( - f"/v1/mcp/server/{server_id}", - headers=self.proxy.transport.master, - json=NoBody(), - response_type=NoBody, - ) - - @step("Sign the key's user in to the MCP server {alias} through the OAuth consent flow") - def seed_user_token(self, alias: str, key: str, storage_state_path: str) -> tuple[str, ...]: - """Drive the interactive authorize dance for `key`'s user so the gateway - stores their upstream token, retried to the shared deadline since the - just-created server and key propagate asynchronously. The LiteLLM key - rides x-litellm-api-key so the gateway binds the token to that user. - Returns the upstream tool names the dance listed, proof the token works.""" - headers = {"x-litellm-api-key": f"Bearer {key}"} - storage = InMemoryTokenStorage() - deadline = time.monotonic() + self.proxy.poll_timeout - last_error: Exception | None = None - while time.monotonic() < deadline: - try: - return asyncio.run(_seed_via_dance(_mcp_url(alias), headers, storage, storage_state_path)) - except Exception as exc: # noqa: BLE001 - retried to the deadline; the last error surfaces below - last_error = exc - time.sleep(self.proxy.poll_interval) - pytest.fail( - f"authorize dance for {alias!r} never completed within {self.proxy.poll_timeout}s; " - f"last error: {last_error!r}" - ) - - @step("List the tools on the MCP server {alias} and call {tool} over the MCP protocol") - def list_and_call( - self, - alias: str, - headers: dict[str, str], - storage: InMemoryTokenStorage, - storage_state_path: str | None, - tool: str, - arguments: dict[str, str], - base_url: str = PROXY_BASE_URL, - identity: Identity | None = None, - allow_upstream_consent: bool = True, - ) -> OauthToolRun: - return asyncio.run( - _list_and_call( - f"{base_url.rstrip('/')}/mcp" if identity is not None else _mcp_url(alias, base_url), - headers, - storage, - storage_state_path, - tool, - arguments, - base_url, - identity, - alias, - allow_upstream_consent, - ) - ) - - @step("List the users with a stored OAuth token for the MCP server") - def server_user_credentials(self, server_id: str) -> tuple[McpServerUserCredentialRow, ...]: - return unwrap( - self.proxy.transport.get( - f"/v1/mcp/server/{server_id}/user-credentials", - headers=self.proxy.transport.master, - params=NoBody(), - response_type=McpServerUserCredentialListResponse, - ) - ).root - - @step("Revoke the user's stored OAuth token for the MCP server") - def revoke_user_token(self, server_id: str, headers: AuthHeaders) -> None: - _ = unwrap( - self.proxy.transport.delete( - f"/v1/mcp/server/{server_id}/oauth-user-credential", - headers=headers, - json=NoBody(), - response_type=McpOauthUserCredentialStatus, - ) - ) - - @step("Send a /chat/completions request to {body.model} with an MCP server attached as a tool") - def chat_with_mcp(self, headers: AuthHeaders, body: ChatBody) -> ChatResponse: - """POST /chat/completions carrying the LiteLLM key in `headers` (either - ingress form) with an MCP server attached in `body.tools`. The gateway - resolves the user from the key and lists/executes the server's tools - with that user's stored upstream token.""" - return unwrap( - self.proxy.transport.post( - "/chat/completions", - headers=headers, - json=body, - response_type=ChatResponse, - ) - ) - - -def build_chat_client(proxy: ProxyClient) -> ChatMcpClient: - return ChatMcpClient(proxy=proxy) diff --git a/tests/e2e/mcp/oauth_gateway.py b/tests/e2e/mcp/oauth_gateway.py deleted file mode 100644 index 2abe8262b6e..00000000000 --- a/tests/e2e/mcp/oauth_gateway.py +++ /dev/null @@ -1,197 +0,0 @@ -"""An owned, source-built OAuth gateway with cold restarts and credential observations. - -Only this child process is restarted. Its database and SSO client survive while -its process-local caches do not; Redis is deliberately absent from its config. -The optional live edge measures headers without recording credentials or bodies. -""" - -from __future__ import annotations - -import os -import subprocess -import sys -import threading -import time -from collections.abc import Callable, Mapping -from contextlib import ExitStack -from dataclasses import dataclass, field -from pathlib import Path -from typing import Final - -import psycopg -from e2e_config import INHERITED_ENV_PREFIXES, available_port -from e2e_http import NoBody -from e2e_metadata import step -from idp import Keycloak, stop_process_group -from proxy_client import ProxyClient, build_proxy_client -from psycopg.rows import class_row -from pydantic import BaseModel, SecretStr, ValidationError - - -class StoredOAuth(BaseModel): - type: str - access_token: SecretStr - - -@dataclass(frozen=True, slots=True) -class CredentialRow: - credential_b64: str = field(repr=False) - - -@step("Read the user's stored OAuth credential for the MCP server from the database and decrypt it") -def stored_oauth(user_id: str, server_id: str) -> StoredOAuth: - """Read the encrypted credential because management APIs omit the plaintext token.""" - from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper - - with psycopg.Connection[CredentialRow].connect( - os.environ["DATABASE_URL"], row_factory=class_row(CredentialRow) - ) as conn: - row: Final = conn.execute( - 'SELECT credential_b64 FROM "LiteLLM_MCPUserCredentials" WHERE user_id = %s AND server_id = %s', - (user_id, server_id), - ).fetchone() - assert row is not None, "canonical user/server has no persisted credential" - plaintext: Final = decrypt_value_helper( - row.credential_b64, "e2e_mcp_oauth", exception_type="debug", return_original_value=False - ) - assert plaintext is not None, "persisted credential must decrypt with the gateway salt" - assert plaintext != row.credential_b64, "persisted credential must be encrypted" - try: - credential: Final = StoredOAuth.model_validate_json(plaintext) - except ValidationError: - raise AssertionError("decrypted credential is not an OAuth payload") from None - assert credential.type == "oauth2" - assert bool(credential.access_token.get_secret_value()), "stored upstream token is empty" - return credential - - -class RpcMethod(BaseModel): - method: str = "" - - -@dataclass(slots=True) -class OAuthObservation: - gateway_token: str = field(default="", repr=False) - _seen: tuple[tuple[str, str, bool], ...] = field(default=(), init=False, repr=False) - _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False) - - def observe(self, url: str, headers: Mapping[str, str], body: bytes | None) -> None: - if body is None or not url.endswith("/mcp"): - return - try: - operation: Final = RpcMethod.model_validate_json(body).method - except ValidationError: - return - if operation not in ("tools/list", "tools/call"): - return - received: Final = headers.get("authorization", "") - gateway_leaked: Final = any( - value in (self.gateway_token, f"Bearer {self.gateway_token}") for value in headers.values() - ) - with self._lock: - self._seen = (*self._seen, (operation, received, gateway_leaked)) - - def assert_forwarded(self, expected: StoredOAuth) -> None: - with self._lock: - snapshot: Final = self._seen - self._seen = () - assert {item[0] for item in snapshot} == {"tools/list", "tools/call"}, "missing upstream observations" - expected_header: Final = f"Bearer {expected.access_token.get_secret_value()}" - assert all(item[1] == expected_header for item in snapshot), "upstream bearer did not match the stored token" - assert all(not item[2] for item in snapshot), "gateway bearer leaked to the upstream" - - -@dataclass(slots=True) -class OAuthGateway: - base_url: str - proxy: ProxyClient - _environment: Mapping[str, str] = field(repr=False) - _command: tuple[str, ...] = field(repr=False) - _log_path: Path - _child: subprocess.Popen[bytes] | None = field(default=None, init=False, repr=False) - - @step("Start the separate LiteLLM proxy and wait for /health/liveliness") - def start(self) -> None: - with self._log_path.open("ab") as log: - self._child = subprocess.Popen( - self._command, - env=self._environment, - stdout=log, - stderr=log, - start_new_session=True, - ) - deadline: Final = time.monotonic() + 120 - while time.monotonic() < deadline: - assert self._child.poll() is None, "owned OAuth gateway exited; inspect its private log" - result = self.proxy.transport.probe("/health/liveliness", params=NoBody()) - if result.status_code == 200: - return - time.sleep(0.5) - raise AssertionError("owned OAuth gateway did not become ready") - - @step("Stop the separate LiteLLM proxy") - def stop(self) -> None: - if self._child is not None: - stop_process_group(self._child) - assert self._child.poll() is not None, "old gateway process is still alive" - - @step("Restart the separate LiteLLM proxy process so its in-memory caches start empty") - def restart(self) -> None: - assert self._child is not None - previous: Final = self._child.pid - self.stop() - self.start() - assert self._child.pid != previous, "gateway restart did not create a new process" - - -@step("Start a separate LiteLLM proxy from source with JWT auth against Keycloak") -def owned_gateway(idp: Keycloak, directory: Path, cleanup: ExitStack) -> OAuthGateway: - for name in ("DATABASE_URL", "LITELLM_LICENSE", "LITELLM_SALT_KEY", "LITELLM_MASTER_KEY"): - assert os.environ.get(name), f"{name} is required for the owned OAuth gateway" - port: Final = available_port() - base_url: Final = f"http://127.0.0.1:{port}" - - def defer(callback: Callable[[], object]) -> None: - cleanup.callback(callback) - - browser: Final = idp.browser_client(callback_url=f"{base_url}/sso/callback", defer=defer) - config: Final = directory / "oauth-gateway.yaml" - config.write_text( - "model_list: []\n" - "general_settings:\n" - " master_key: os.environ/LITELLM_MASTER_KEY\n" - " database_url: os.environ/DATABASE_URL\n" - " enable_jwt_auth: true\n" - " litellm_jwtauth:\n" - " user_id_jwt_field: sub\n" - " user_email_jwt_field: email\n" - " team_ids_jwt_field: groups\n" - " user_id_upsert: true\n" - ) - environment: Final = { - **{key: value for key, value in os.environ.items() if not key.startswith(INHERITED_ENV_PREFIXES)}, - **browser.environment(idp.discovery()), - "PROXY_BASE_URL": base_url, - "JWT_PUBLIC_KEY_URL": idp.jwks_url, - "JWT_ISSUER": idp.issuer, - "JWT_AUDIENCE": "litellm-e2e", - "DISABLE_SCHEMA_UPDATE": "true", - "STORE_MODEL_IN_DB": "True", - "PYTHONPATH": str(Path(__file__).resolve().parents[3]), - } - gateway: Final = OAuthGateway( - base_url=base_url, - proxy=build_proxy_client( - base_url=base_url, - control_plane_base_url=base_url, - replica_urls=(base_url,), - control_replica_urls=(base_url,), - master_key=os.environ["LITELLM_MASTER_KEY"], - ), - _environment=environment, - _command=(sys.executable, "-m", "litellm.proxy.proxy_cli", "--config", str(config), "--port", str(port)), - _log_path=directory / "oauth-gateway.log", - ) - cleanup.callback(gateway.stop) - gateway.start() - return gateway diff --git a/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py b/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py deleted file mode 100644 index 9d7e4713963..00000000000 --- a/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py +++ /dev/null @@ -1,218 +0,0 @@ -"""On-demand e2e: a chat completion drives a gateway-managed OAuth MCP server. - -The real end-user flow for MCP over an OAuth server: a user registers a Linear -authorization_code server, authorizes it once so the gateway stores their -upstream token, then sends a normal /chat/completions request with the Linear -MCP attached. The gateway resolves the user from the LiteLLM key, lists Linear's -tools with the stored per-user token, lets the model call one, executes it -upstream with that token, and returns the answer. This is proven against the -real Linear MCP server (mcp.linear.app) and a real Anthropic model, once per -documented ingress header (x-litellm-api-key and Authorization). - -The authorize dance is seeded through the mcp SDK's OAuthClientProvider; the one -step Linear cannot auto-approve is the human consent, so it is captured once out -of band (mcp/linear_session_capture.py) into a saved browser session and a -headless Chromium clicks Approve every run. The test therefore skips unless -E2E_LINEAR_STORAGE_STATE points at that session, so it never runs on the per-PR -CI path; it is a nightly/on-demand real-server smoke test. - -Fail-before-fix: without the stored per-user token the gateway lists no Linear -tools, so mcp_list_tools comes back empty, nothing is called, and the -assertions fail; a served, called, non-empty Linear tool proves the gateway -pulled and used the user's token. -""" - -from __future__ import annotations - -import os - -import pytest - -from e2e_config import CHEAP_ANTHROPIC_MODEL, LINEAR_MCP_URL, LINEAR_STORAGE_STATE, unique_marker -from e2e_http import AuthHeaders -from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta -from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, KeyGenerateBody, McpChatTool, McpServerCreateBody, ObjectPermission -from proxy_client import ProxyClient - -pytest.importorskip("mcp", reason="mcp SDK not installed; run `uv sync --inexact --group e2e-dev`") -pytest.importorskip( - "playwright.async_api", - reason="playwright not installed; run `uv pip install playwright` and `playwright install chromium`", -) - -from oauth_chat_client import ChatMcpClient, build_chat_client # noqa: E402 # imports follow the importorskip guards - -pytestmark = [ - pytest.mark.e2e, - pytest.mark.skipif( - not LINEAR_STORAGE_STATE or not os.path.exists(LINEAR_STORAGE_STATE), - reason="set E2E_LINEAR_STORAGE_STATE to a Linear session captured via mcp/linear_session_capture.py", - ), -] - -# Pinned from a live dance during verification (never guessed); the gateway -# prefixes every upstream tool name with the server alias. list_teams is a -# read-only Linear tool that takes no arguments and returns the caller's teams. -LINEAR_READONLY_TOOL = "list_teams" -LINEAR_PROMPT = "Use the list_teams tool to list my Linear teams, then reply with the name of one of them." - - -@pytest.fixture(scope="session") -def chat_client(proxy: ProxyClient) -> ChatMcpClient: - return build_chat_client(proxy) - - -class TestMcpChatCompletionOauth: - """A scoped internal-user key on a real Linear authorization_code server, - used through /chat/completions once per ingress header: the gateway pulls - the user's stored upstream token, lists and executes Linear's tools during - the completion, and returns the answer.""" - - @pytest.mark.covers("mcp.list_tools.oauth.succeeds") - @pytest.mark.covers("mcp.call_tool.oauth.succeeds") - @meta( - Subject( - domain=Domain.MCP, - route=Route.CHAT_COMPLETIONS, - providers=(Provider.ANTHROPIC,), - models=(CHEAP_ANTHROPIC_MODEL,), - capabilities=(Capability.FUNCTION_CALLING,), - mode=Mode.NONSTREAM, - ) - ) - def test_chat_completion_uses_linear_with_x_litellm_api_key_header( - self, chat_client: ChatMcpClient, resources: ResourceManager - ) -> None: - marker = unique_marker() - alias = f"e2elinear{marker}" - created = chat_client.create_server( - McpServerCreateBody( - alias=alias, - url=LINEAR_MCP_URL, - allow_all_keys=False, - auth_type="oauth2", - oauth2_flow="authorization_code", - ) - ) - resources.defer(lambda: chat_client.delete_server(created.server_id)) - - stored = chat_client.server_info(created.server_id) - assert stored.auth_type == "oauth2" - assert stored.oauth2_flow == "authorization_code" - assert stored.allow_all_keys is False - - key = chat_client.proxy.generate_key( - KeyGenerateBody( - user_id="e2e-test-user", - object_permission=ObjectPermission(mcp_servers=[created.server_id]), - ) - ) - resources.defer(lambda: chat_client.proxy.delete_key(key)) - - seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE) - assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, ( - f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}" - ) - - response = chat_client.chat_with_mcp( - AuthHeaders.model_validate({"x-litellm-api-key": f"Bearer {key}"}), - ChatBody( - model=CHEAP_ANTHROPIC_MODEL, - messages=[ChatMessage(role="user", content=LINEAR_PROMPT)], - tools=[ - McpChatTool( - server_url=f"litellm_proxy/mcp/{alias}", - server_label=alias, - require_approval="never", - ) - ], - ), - ) - - message = response.choices[0].message - assert message is not None and message.content, f"completion returned no answer: {response}" - meta = message.provider_specific_fields - assert meta is not None, f"no MCP metadata on the completion: {response}" - listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function} - assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, ( - f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}" - ) - results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"] - assert results and results[0].result, ( - f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}" - ) - - @pytest.mark.covers("mcp.list_tools.oauth.succeeds") - @pytest.mark.covers("mcp.call_tool.oauth.succeeds") - @meta( - Subject( - domain=Domain.MCP, - route=Route.CHAT_COMPLETIONS, - providers=(Provider.ANTHROPIC,), - models=(CHEAP_ANTHROPIC_MODEL,), - capabilities=(Capability.FUNCTION_CALLING,), - mode=Mode.NONSTREAM, - ) - ) - def test_chat_completion_uses_linear_with_authorization_bearer_header( - self, chat_client: ChatMcpClient, resources: ResourceManager - ) -> None: - marker = unique_marker() - alias = f"e2elinear{marker}" - created = chat_client.create_server( - McpServerCreateBody( - alias=alias, - url=LINEAR_MCP_URL, - allow_all_keys=False, - auth_type="oauth2", - oauth2_flow="authorization_code", - ) - ) - resources.defer(lambda: chat_client.delete_server(created.server_id)) - - stored = chat_client.server_info(created.server_id) - assert stored.auth_type == "oauth2" - assert stored.oauth2_flow == "authorization_code" - assert stored.allow_all_keys is False - - key = chat_client.proxy.generate_key( - KeyGenerateBody( - user_id="e2e-test-user", - object_permission=ObjectPermission(mcp_servers=[created.server_id]), - ) - ) - resources.defer(lambda: chat_client.proxy.delete_key(key)) - - seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE) - assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, ( - f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}" - ) - - response = chat_client.chat_with_mcp( - AuthHeaders.model_validate({"authorization": f"Bearer {key}"}), - ChatBody( - model=CHEAP_ANTHROPIC_MODEL, - messages=[ChatMessage(role="user", content=LINEAR_PROMPT)], - tools=[ - McpChatTool( - server_url=f"litellm_proxy/mcp/{alias}", - server_label=alias, - require_approval="never", - ) - ], - ), - ) - - message = response.choices[0].message - assert message is not None and message.content, f"completion returned no answer: {response}" - meta = message.provider_specific_fields - assert meta is not None, f"no MCP metadata on the completion: {response}" - listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function} - assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, ( - f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}" - ) - results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"] - assert results and results[0].result, ( - f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}" - ) diff --git a/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py b/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py deleted file mode 100644 index b57c941cb18..00000000000 --- a/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py +++ /dev/null @@ -1,216 +0,0 @@ -"""Real OAuth consent, immediate MCP operations and cold-restart persistence. - -Aggregate SSO uses the SDK's normal authentication. The per-server variant is -explicitly a configured two-header client, not an Authorization-only OAuth host. -The observed variants forward to the same real Linear upstream and compare its -bearer at the forwarding boundary; direct variants retain unmodified discovery. -""" - -from __future__ import annotations - -import os -from collections.abc import Iterator -from contextlib import ExitStack -from pathlib import Path -from types import MappingProxyType -from typing import Final, Literal - -import pytest -from e2e_config import LINEAR_MCP_URL, LINEAR_READONLY_TOOL, LINEAR_STORAGE_STATE, unique_marker -from e2e_http import AuthHeaders, NoBody, get_external, unwrap -from e2e_metadata import Domain, Route, Subject, meta -from idp import Identity, Keycloak -from lifecycle import ResourceManager -from models import ( - McpOauthCredentials, - McpServerCreateBody, - ObjectPermission, - TeamMemberAddBody, - TeamMemberEntry, - TeamUpdateBody, -) -from oauth_chat_client import ChatMcpClient, InMemoryTokenStorage, OauthToolRun, build_chat_client -from oauth_gateway import OAuthGateway, OAuthObservation, owned_gateway, stored_oauth -from provider_edge import LiveEdge, start_provider_edge -from proxy_client import ProxyClient -from pydantic import BaseModel, ValidationError - -pytestmark = [pytest.mark.e2e, pytest.mark.mcp_oauth_live, pytest.mark.provider_live] - - -class OAuthMetadata(BaseModel): - issuer: str - authorization_endpoint: str - token_endpoint: str - registration_endpoint: str - - -class LinearTeam(BaseModel): - id: str - name: str - - -class LinearTeams(BaseModel): - teams: tuple[LinearTeam, ...] - - -def assert_tool_result(run: OauthToolRun, tool: str) -> None: - assert tool in run.tools - assert run.is_error is False - try: - result: Final = LinearTeams.model_validate_json(run.text) - except ValidationError: - raise AssertionError("list_teams did not return the expected teams payload") from None - assert result.teams, "the test workspace must contain at least one team" - assert all(team.id and team.name for team in result.teams), "team results must contain identifiers and names" - - -@pytest.fixture(scope="module") -def oauth_gateway(idp: Keycloak, tmp_path_factory: pytest.TempPathFactory) -> Iterator[OAuthGateway]: - assert LINEAR_STORAGE_STATE and Path(LINEAR_STORAGE_STATE).is_file(), ( - "E2E_LINEAR_STORAGE_STATE must name a captured Linear login; see mcp/linear_session_capture.py" - ) - assert os.environ.get("E2E_FIXTURE_MODE", "live") == "live", "OAuth acceptance cannot use replay" - with ExitStack() as cleanup: - yield owned_gateway(idp, tmp_path_factory.mktemp("mcp-oauth"), cleanup) - - -@pytest.fixture(scope="module") -def proxy(oauth_gateway: OAuthGateway) -> ProxyClient: - return oauth_gateway.proxy - - -@pytest.fixture(scope="module") -def client(proxy: ProxyClient) -> ChatMcpClient: - return build_chat_client(proxy) - - -class TestMcpOauthHappyPath: - @pytest.mark.covers("mcp.list_tools.oauth.succeeds") - @pytest.mark.covers("mcp.call_tool.oauth.succeeds") - @pytest.mark.covers("mcp.call_tool.oauth.persists_across_processes") - @meta( - Subject( - domain=Domain.MCP, - route=Route.MCP, - ) - ) - @pytest.mark.parametrize("route", ("aggregate_sso", "explicit_header_jwt")) - @pytest.mark.parametrize("observed", (False, True), ids=("direct", "observed")) - def test_consent_list_call_and_cold_restart( - self, - client: ChatMcpClient, - resources: ResourceManager, - jwt_identity: Identity, - idp: Keycloak, - oauth_gateway: OAuthGateway, - route: Literal["aggregate_sso", "explicit_header_jwt"], - observed: bool, - ) -> None: - alias: Final = f"e2elinear{unique_marker()}" - tool: Final = f"{alias}-{LINEAR_READONLY_TOOL}" - token: Final = idp.access_token(jwt_identity) - observation: Final = OAuthObservation(gateway_token=token) - edge: Final = ( - start_provider_edge( - LiveEdge(observe_request=observation.observe), - mounts=MappingProxyType( - {"linear": "https://mcp.linear.app", ".well-known": "https://mcp.linear.app/.well-known"} - ), - ) - if observed - else None - ) - if edge is not None: - resources.defer(edge.shutdown) - metadata: Final = ( - unwrap( - get_external( - "https://mcp.linear.app/.well-known/oauth-authorization-server", - response_type=OAuthMetadata, - ) - ) - if observed - else None - ) - created: Final = client.create_server( - McpServerCreateBody( - alias=alias, - server_name=alias, - url=f"{edge.edge.api_base('linear')}/mcp" if edge is not None else LINEAR_MCP_URL, - transport="http", - allow_all_keys=False, - auth_type="oauth2", - oauth2_flow="authorization_code", - per_server_oauth_discovery=route == "explicit_header_jwt", - issuer=metadata.issuer if metadata else None, - authorization_url=metadata.authorization_endpoint if metadata else None, - token_url=metadata.token_endpoint if metadata else None, - registration_url=metadata.registration_endpoint if metadata else None, - credentials=McpOauthCredentials(upstream_resource=LINEAR_MCP_URL) if observed else None, - ) - ) - resources.defer(lambda: client.delete_server(created.server_id)) - assert client.server_user_credentials(created.server_id) == (), ( - "scenario must start without upstream credentials" - ) - unwrap( - client.proxy.transport.post( - "/team/member_add", - headers=client.proxy.transport.master, - json=TeamMemberAddBody( - team_id=jwt_identity.group, member=TeamMemberEntry(user_id=jwt_identity.user_id, role="user") - ), - response_type=NoBody, - ) - ) - client.proxy.update_team( - TeamUpdateBody( - team_id=jwt_identity.group, - object_permission=ObjectPermission(mcp_servers=[created.server_id]), - ) - ) - headers: Final = {"x-litellm-api-key": f"Bearer {token}"} if route == "explicit_header_jwt" else {} - resources.defer( - lambda: client.revoke_user_token( - created.server_id, - AuthHeaders(authorization=f"Bearer {idp.access_token(jwt_identity)}"), - ) - ) - identity: Final = jwt_identity if route == "aggregate_sso" else None - first: Final = client.list_and_call( - alias, - headers, - InMemoryTokenStorage(), - LINEAR_STORAGE_STATE, - tool, - {}, - base_url=oauth_gateway.base_url, - identity=identity, - ) - assert_tool_result(first, tool) - credentials: Final = client.server_user_credentials(created.server_id) - assert len(credentials) == 1 - assert credentials[0].user_id == jwt_identity.user_id - assert credentials[0].credential_type == "oauth2" - first_stored_oauth: Final = stored_oauth(jwt_identity.user_id, created.server_id) - if observed: - observation.assert_forwarded(first_stored_oauth) - oauth_gateway.restart() - fresh_token: Final = idp.access_token(jwt_identity) - observation.gateway_token = fresh_token - second: Final = client.list_and_call( - alias, - {"Authorization": f"Bearer {fresh_token}"} if identity is None else {}, - InMemoryTokenStorage(), - LINEAR_STORAGE_STATE if identity is not None else None, - tool, - {}, - base_url=oauth_gateway.base_url, - identity=identity, - allow_upstream_consent=False, - ) - assert_tool_result(second, tool) - second_stored_oauth: Final = stored_oauth(jwt_identity.user_id, created.server_id) - if observed: - observation.assert_forwarded(second_stored_oauth) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index d42518650fb..438c1f1fce6 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -730,28 +730,14 @@ class McpInfo(BaseModel): logo_url: str | None = None -class McpOauthCredentials(BaseModel): - upstream_resource: str - - class McpServerCreateBody(BaseModel): - """POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is - `oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints - are discovered and registered via DCR when left unset. `allow_all_keys` - false scopes the server to keys granted it through object_permission.""" + """POST /v1/mcp/server. `allow_all_keys` false scopes the server to keys + granted it through object_permission.""" alias: str url: str transport: str = "http" allow_all_keys: bool = True - auth_type: str | None = None - oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None - per_server_oauth_discovery: bool | None = None - issuer: str | None = None - authorization_url: str | None = None - token_url: str | None = None - registration_url: str | None = None - credentials: McpOauthCredentials | None = Field(default=None, repr=False) server_name: str | None = None description: str | None = None mcp_info: McpInfo | None = None @@ -791,26 +777,6 @@ class McpServerListResponse(RootModel[list[McpServerRow]]): """GET /v1/mcp/server answers with a bare array of servers.""" -class McpServerUserCredentialRow(BaseModel): - user_id: str - credential_type: Literal["oauth2", "byok"] - expires_at: str | None = None - connected_at: str | None = None - updated_at: str - - -class McpServerUserCredentialListResponse(RootModel[tuple[McpServerUserCredentialRow, ...]]): - """GET /v1/mcp/server/{server_id}/user-credentials answers with a bare array.""" - - -class McpOauthUserCredentialStatus(BaseModel): - server_id: str - has_credential: bool - expires_at: str | None = None - is_expired: bool = False - connected_at: str | None = None - - class ToolsetTool(BaseModel): server_id: str tool_name: str diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index eab59c3c77c..891c5fc282c 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -13,7 +13,6 @@ markers = prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set - mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless E2E_MCP_OAUTH_LIVE is set provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set owned_gateway: boots its own proxy from source against the stack's Postgres, so it needs DATABASE_URL on the pytest host; deselected unless E2E_OWNED_GATEWAY is set otel_v2: needs a proxy running with LITELLM_OTEL_V2=true; deselected unless E2E_OTEL_V2 is set diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index 6f59f694352..563d35bf794 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -697,6 +697,54 @@ def test_per_user_authorization_code_with_pkce_binds_the_token_to_the_authorizin ) +@pytest.mark.timeout(600) +def test_per_user_authorization_code_token_is_read_back_from_the_database_by_a_fresh_proxy_without_redis( + gateway: Gateway, tmp_path: Path +) -> None: + with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario: + alias: Final = "cold" + uuid.uuid4().hex[:8] + identity: Final = _register_oauth( + scenario, + peer, + auth, + alias, + auth_type="oauth2", + oauth2_flow="authorization_code", + credentials={"client_id": "cold-client", "client_secret": "cold-secret", "scopes": ["tools.call"]}, + ) + owner_id: Final = scenario.user() + owner: Final = scenario.key(user_id=owner_id, object_permission={"mcp_servers": [identity]}) + registered: Final = gateway.client.post( + f"/{alias}/register", json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"} + ) + assert registered.status_code in (200, 201), registered.text + client_id: Final = str(_response_object(registered)["client_id"]) + pkce: Final = _Pkce(secrets.token_urlsafe(32)) + code, _ = _authorize_through_gateway(gateway, auth, alias, owner, client_id, pkce) + redeemed: Final = _redeem(gateway, alias, owner, client_id, code, pkce) + assert redeemed.status_code == 200, redeemed.text + token: Final = _issued_token(redeemed.json()) + stored: Final = read_rows( + 'SELECT credential_b64 FROM "LiteLLM_MCPUserCredentials" WHERE user_id = %s AND server_id = %s', + (owner_id, identity), + ) + assert len(stored) == 1, stored + stored_bytes: Final = base64.urlsafe_b64decode(str(stored[0]["credential_b64"])) + assert token.encode() not in stored_bytes, "stored token is plaintext" + config: Final = tmp_path / "no_redis.yaml" + config.write_text( + "general_settings:\n" + " master_key: os.environ/LITELLM_MASTER_KEY\n" + " database_url: os.environ/DATABASE_URL\n" + " store_model_in_db: true\n" + ) + peer.drain() + with owned_proxy(gateway, tmp_path, {}, config=config, remove_environment=("REDIS_HOST", "REDIS_PORT")) as cold: + called: Final = call_tool(cold, owner, identity, f"{alias}-add", ADD) + assert called.status_code == 200, called.text + assert _authorizations(peer) == (f"Bearer {token}".encode(),) + + def test_authorization_request_without_pkce_is_refused_before_reaching_the_authorization_server( gateway: Gateway, ) -> None: diff --git a/tests/unit/test_circleci_path_filter.py b/tests/unit/test_circleci_path_filter.py index 0856115e39b..f94b2fa863a 100644 --- a/tests/unit/test_circleci_path_filter.py +++ b/tests/unit/test_circleci_path_filter.py @@ -52,7 +52,7 @@ CI = [".github/workflows/test-unit.yml"] ("mcp-dependencies", ["pyproject.toml"], "run"), ("mcp-dependencies", ["uv.lock"], "run"), ("mcp-dependencies", ["litellm/experimental_mcp_client/client.py"], "run"), - ("mcp-dependencies", ["tests/e2e/mcp/oauth_chat_client.py"], "run"), + ("mcp-dependencies", ["tests/e2e/mcp/datadog_mcp.py"], "run"), ("mcp-dependencies", ["litellm-proxy-extras/pyproject.toml"], "run"), ("mcp-dependencies", ["scripts/check_mcp_sdk_install.py"], "run"), ("mcp-dependencies", [".github/workflows/test-mcp-dependency-resolution.yml"], "run"), From 847380114532d990872819ed896d89cc42806585 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:28:53 -0700 Subject: [PATCH 15/41] ci: remove the Together AI and weekly cost map sync workflows (#48473) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit 57ac10731149fbc2a5612a811c1dc15f6780dac1) --- ...to_update_price_and_context_window_file.py | 319 ---------- .../auto_update_price_and_context_window.yml | 39 -- .github/workflows/sync-together-ai-models.yml | 70 --- scripts/sync_together_ai_models.py | 547 ------------------ ...to_update_price_and_context_window_file.py | 222 ------- tests/unit/test_sync_together_ai_models.py | 378 ------------ 6 files changed, 1575 deletions(-) delete mode 100644 .github/scripts/auto_update_price_and_context_window_file.py delete mode 100644 .github/workflows/auto_update_price_and_context_window.yml delete mode 100644 .github/workflows/sync-together-ai-models.yml delete mode 100644 scripts/sync_together_ai_models.py delete mode 100644 tests/unit/test_auto_update_price_and_context_window_file.py delete mode 100644 tests/unit/test_sync_together_ai_models.py diff --git a/.github/scripts/auto_update_price_and_context_window_file.py b/.github/scripts/auto_update_price_and_context_window_file.py deleted file mode 100644 index 4c9a05bd4f0..00000000000 --- a/.github/scripts/auto_update_price_and_context_window_file.py +++ /dev/null @@ -1,319 +0,0 @@ -import asyncio -import aiohttp -import json -import math -from typing import Any - -# Asynchronously fetch data from a given URL -async def fetch_data(url): - try: - # Create an asynchronous session - async with aiohttp.ClientSession() as session: - # Send a GET request to the URL - async with session.get(url) as resp: - # Raise an error if the response status is not OK - resp.raise_for_status() - # Parse the response JSON - resp_json = await resp.json() - print("Fetch the data from URL.") - # Return the 'data' field from the JSON response - return resp_json['data'] - except Exception as e: - # Print an error message if fetching data fails - print("Error fetching data from URL:", e) - return None - - -FRIENDLI_API_URL = "https://api.friendli.ai/serverless/v1/models" -FRIENDLI_PROVIDER = "friendliai" - -INHERITABLE_BASE_KEYS = ( - "supports_pdf_input", - "supports_assistant_prefill", - "supports_adaptive_thinking", - "supports_output_config", -) - -REASONING_EFFORT_LEVEL_ORDER = ("none", "minimal", "low", "medium", "high", "xhigh", "max") - - -def _find_base_model_entry(base_model: str, local_data: dict) -> str | None: - if not base_model: - return None - bm_tail = base_model.split("/")[-1].lower() - if base_model in local_data: - return base_model - for key in local_data: - if key.startswith("sample_spec") or key == "fallback_generalizations": - continue - if key.split("/")[-1].lower() == bm_tail: - return key - return None - - -def _reasoning_effort_levels(reasoning_options: list) -> list: - offered = { - val - for opt in reasoning_options or [] - if opt.get("type") == "effort" - for val in opt.get("values", []) - } - return [level for level in REASONING_EFFORT_LEVEL_ORDER if level in offered] - - -def _valid_token_price(value: object) -> bool: - try: - price = float(value) # pyright: ignore[reportArgumentType] # non-numeric values are rejected via the except - except (TypeError, ValueError): - return False - return math.isfinite(price) and price >= 0 - - -def _has_valid_token_prices(pricing: dict | None) -> bool: - prices = pricing or {} - return _valid_token_price(prices.get("input")) and _valid_token_price(prices.get("output")) - - -def _pricing(pricing: dict) -> dict: - out: dict[str, Any] = {} - if not pricing: - return out - if "input" in pricing: - out["input_cost_per_token"] = float(pricing["input"]) - if "output" in pricing: - out["output_cost_per_token"] = float(pricing["output"]) - if "input_cache_read" in pricing and pricing["input_cache_read"] is not None: - out["cache_read_input_token_cost"] = float(pricing["input_cache_read"]) - return out - - -def _modality_flags(input_mods: list) -> dict: - mods = input_mods or [] - has_image = "image" in mods - return { - "supports_vision": has_image, - "supports_image_input": has_image, - "supports_video_input": "video" in mods, - } - - -def transform_friendli_data(data: list, local_data: dict) -> dict: - transformed: dict[str, dict] = {} - if not data: - return transformed - for model in data: - # An unpriced row must never wholesale-replace an already priced local entry: - # missing prices cost-calculate as zero, silently zeroing tracked spend - if not _has_valid_token_prices(model.get("pricing")): - continue - model_id = model["id"] - base_model = model.get("base_model") or "" - entry: dict[str, Any] = { - "litellm_provider": FRIENDLI_PROVIDER, - } - - base_key = _find_base_model_entry(base_model, local_data) - if base_key: - base_entry = local_data[base_key] - for k in INHERITABLE_BASE_KEYS: - if k in base_entry: - entry[k] = base_entry[k] - - ctx = model.get("context_length") - if ctx is not None: - entry["max_input_tokens"] = int(ctx) - max_out = model.get("max_completion_tokens") - if max_out is not None: - entry["max_output_tokens"] = int(max_out) - entry["max_tokens"] = int(max_out) - - pricing = _pricing(model.get("pricing", {})) - entry.update(pricing) - entry["supports_prompt_caching"] = "cache_read_input_token_cost" in pricing - - reasoning = model.get("reasoning") is True - entry["supports_reasoning"] = reasoning - if reasoning: - entry["reasoning_effort_levels"] = _reasoning_effort_levels( - model.get("reasoning_options", []) - ) - - func = model.get("functionality", {}) - entry["supports_function_calling"] = func.get("tool_call") is True - entry["supports_parallel_function_calling"] = func.get("parallel_tool_call") is True - is_struct = func.get("structured_output") is True - entry["supports_response_schema"] = is_struct - entry["supports_native_structured_output"] = is_struct - entry["supports_system_messages"] = func.get("system_messages") is True - entry["supports_tool_choice"] = func.get("tool_choice") is True - - entry.update(_modality_flags(model.get("input_modalities", []))) - - entry["mode"] = model.get("mode", "chat") - - desc = model.get("description") - if desc: - entry["comment"] = desc - - dep = model.get("deprecation_date") - if dep: - entry["deprecation_date"] = dep.split("T")[0] - - entry["source"] = FRIENDLI_API_URL - - transformed[f"{FRIENDLI_PROVIDER}/{model_id}"] = entry - return transformed - -# Synchronize local data with remote data -def sync_local_data_with_remote(local_data, remote_data, replace_keys=frozenset()): - # Update existing keys in local_data with values from remote_data - # (replace_keys entries are swapped wholesale so a field the remote catalog - # dropped, e.g. cache pricing, cannot survive as a stale value) - for key in (set(local_data) & set(remote_data)): - if key in replace_keys: - local_data[key] = remote_data[key] - else: - local_data[key].update(remote_data[key]) - - # Add new keys from remote_data to local_data - for key in (set(remote_data) - set(local_data)): - local_data[key] = remote_data[key] - -# Write data to the json file -def write_to_file(file_path, data): - try: - # Open the file in write mode - with open(file_path, "w") as file: - # Dump the data as JSON into the file - json.dump(data, file, indent=4) - print("Values updated successfully.") - except Exception as e: - # Print an error message if writing to file fails - print("Error updating JSON file:", e) - -# Update the existing models and add the missing models for OpenRouter -def transform_openrouter_data(data): - transformed = {} - if not data: - return transformed - for row in data: - # Add the fields 'max_tokens' and 'input_cost_per_token' - obj = { - "max_tokens": row["context_length"], - "input_cost_per_token": float(row["pricing"]["prompt"]), - } - - # Add 'max_output_tokens' as a field if it is not None - if "top_provider" in row and "max_completion_tokens" in row["top_provider"] and row["top_provider"]["max_completion_tokens"] is not None: - obj['max_output_tokens'] = int(row["top_provider"]["max_completion_tokens"]) - - # Add the field 'output_cost_per_token' - obj.update({ - "output_cost_per_token": float(row["pricing"]["completion"]), - }) - - # Add field 'input_cost_per_image' if it exists and is non-zero - if "pricing" in row and "image" in row["pricing"] and float(row["pricing"]["image"]) != 0.0: - obj['input_cost_per_image'] = float(row["pricing"]["image"]) - - # Add the fields 'litellm_provider' and 'mode' - obj.update({ - "litellm_provider": "openrouter", - "mode": "chat" - }) - - # Add the 'supports_vision' field if the modality is 'multimodal' - if row.get('architecture', {}).get('modality') == 'multimodal': - obj['supports_vision'] = True - - # Use a composite key to store the transformed object - transformed[f'openrouter/{row["id"]}'] = obj - - return transformed - -# Update the existing models and add the missing models for Vercel AI Gateway -def transform_vercel_ai_gateway_data(data): - transformed = {} - if not data: - return transformed - for row in data: - # Rows without token pricing or token limits (video/embedding models) previously KeyError'd the whole sync - if any(row.get(k) is None for k in ("context_window", "max_tokens")) or any( - row.get("pricing", {}).get(k) is None for k in ("input", "output") - ): - continue - obj = { - "max_tokens": row["context_window"], - "input_cost_per_token": float(row["pricing"]["input"]), - "output_cost_per_token": float(row["pricing"]["output"]), - 'max_output_tokens': row['max_tokens'], - 'max_input_tokens': row["context_window"], - } - - # Handle cache pricing if available - if "pricing" in row: - if "input_cache_read" in row["pricing"] and row["pricing"]["input_cache_read"] is not None: - obj['cache_read_input_token_cost'] = float(f"{float(row['pricing']['input_cache_read']):e}") - - if "input_cache_write" in row["pricing"] and row["pricing"]["input_cache_write"] is not None: - obj['cache_creation_input_token_cost'] = float(f"{float(row['pricing']['input_cache_write']):e}") - - mode = "embedding" if "embedding" in row["id"].lower() else "chat" - - obj.update({"litellm_provider": "vercel_ai_gateway", "mode": mode}) - - transformed[f'vercel_ai_gateway/{row["id"]}'] = obj - - return transformed - - -# Load local data from a specified file -def load_local_data(file_path): - try: - # Open the file in read mode - with open(file_path, "r") as file: - # Load and return the JSON data - return json.load(file) - except FileNotFoundError: - # Print an error message if the file is not found - print("File not found:", file_path) - return None - except json.JSONDecodeError as e: - # Print an error message if JSON decoding fails - print("Error decoding JSON:", e) - return None - -def main(): - local_file_path = "model_prices_and_context_window.json" # Path to the local data file - openrouter_url = "https://openrouter.ai/api/v1/models" # URL to fetch OpenRouter data - vercel_ai_gateway_url = "https://ai-gateway.vercel.sh/v1/models" # URL to fetch Vercel AI Gateway data - - # Load local data from file - local_data = load_local_data(local_file_path) - - # Fetch OpenRouter data - openrouter_data = asyncio.run(fetch_data(openrouter_url)) - # Transform the fetched OpenRouter data - openrouter_data = transform_openrouter_data(openrouter_data) - - # Fetch Vercel AI Gateway data - vercel_data = asyncio.run(fetch_data(vercel_ai_gateway_url)) - # Transform the fetched Vercel AI Gateway data - vercel_data = transform_vercel_ai_gateway_data(vercel_data) - - friendli_data = asyncio.run(fetch_data(FRIENDLI_API_URL)) - friendli_data = transform_friendli_data(friendli_data, local_data) - - # Combine both datasets - all_remote_data = {**openrouter_data, **vercel_data, **friendli_data} - - # If both local and openrouter data are available, synchronize and save - if local_data and all_remote_data: - sync_local_data_with_remote(local_data, all_remote_data, replace_keys=frozenset(friendli_data)) - write_to_file(local_file_path, local_data) - else: - print("Failed to fetch model data from either local file or URL.") - -# Entry point of the script -if __name__ == "__main__": - main() diff --git a/.github/workflows/auto_update_price_and_context_window.yml b/.github/workflows/auto_update_price_and_context_window.yml deleted file mode 100644 index 7e40a860ee9..00000000000 --- a/.github/workflows/auto_update_price_and_context_window.yml +++ /dev/null @@ -1,39 +0,0 @@ -name: Updates model_prices_and_context_window.json and Create Pull Request - -on: - schedule: - - cron: "0 0 * * 0" # Run every Sundays at midnight - #- cron: "0 0 * * *" # Run daily at midnight - -permissions: - contents: write - pull-requests: write - -jobs: - auto_update_price_and_context_window: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - name: Update JSON Data - run: | - uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py" - - name: Regenerate JSON Schema - run: | - uv run --frozen python ci_cd/generate_model_prices_schema.py - - name: Create Pull Request - run: | - git add model_prices_and_context_window.json model_prices_and_context_window.schema.json - git commit -m "Update model_prices_and_context_window.json file: $(date +'%Y-%m-%d')" - gh pr create --title "Update model_prices_and_context_window.json file" \ - --body "Automated update for model_prices_and_context_window.json" \ - --head auto-update-price-and-context-window-$(date +'%Y-%m-%d') \ - --base main - env: - GH_TOKEN: ${{ secrets.GH_TOKEN }} diff --git a/.github/workflows/sync-together-ai-models.yml b/.github/workflows/sync-together-ai-models.yml deleted file mode 100644 index 18d5e3de1eb..00000000000 --- a/.github/workflows/sync-together-ai-models.yml +++ /dev/null @@ -1,70 +0,0 @@ -name: Sync Together AI model registry - -on: - schedule: - - cron: "30 6 * * *" - workflow_dispatch: - -permissions: - contents: write - pull-requests: write - -jobs: - sync_together_ai_models: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - env: - BASE_BRANCH: ${{ github.event.repository.default_branch }} - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - ref: ${{ env.BASE_BRANCH }} - persist-credentials: false - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - name: Look for an already-open sync PR - id: existing - run: | - open_pr="$(gh pr list --repo "$GITHUB_REPOSITORY" --state open --limit 1000 --json headRefName \ - --jq '[.[].headRefName | select(startswith("litellm_together_registry_sync_"))] | first // empty')" - echo "open_pr=$open_pr" >> "$GITHUB_OUTPUT" - if [ -n "$open_pr" ]; then - echo "An open sync PR already exists on branch $open_pr; skipping this run." - fi - env: - GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }} - - name: Run the sync - if: steps.existing.outputs.open_pr == '' - run: | - uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md" - env: - TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }} - - name: Regenerate the JSON schema - if: steps.existing.outputs.open_pr == '' - run: | - uv run --frozen python ci_cd/generate_model_prices_schema.py - - name: Create a pull request when the registry changed - if: steps.existing.outputs.open_pr == '' - run: | - if git diff --quiet; then - echo "Registry already in sync; no PR needed." - exit 0 - fi - branch="litellm_together_registry_sync_$(date +'%Y-%m-%d')" - git config user.name "github-actions[bot]" - git config user.email "41898282+github-actions[bot]@users.noreply.github.com" - git checkout -b "$branch" - git add model_prices_and_context_window.json \ - litellm/model_prices_and_context_window_backup.json \ - model_prices_and_context_window.schema.json - git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')" - gh auth setup-git - git push origin "$branch" - gh pr create --title "feat(models): sync together_ai model registry" \ - --body-file "$RUNNER_TEMP/pr_body.md" \ - --head "$branch" \ - --base "$BASE_BRANCH" - env: - GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }} diff --git a/scripts/sync_together_ai_models.py b/scripts/sync_together_ai_models.py deleted file mode 100644 index 12b128890f1..00000000000 --- a/scripts/sync_together_ai_models.py +++ /dev/null @@ -1,547 +0,0 @@ -"""Sync the together_ai entries of model_prices_and_context_window.json with Together's live serverless catalog. - -Pulls ``GET https://api.together.ai/v1/models?serverless`` plus the deprecations doc, maps API fields onto -registry fields, merges the reviewed capability rules below for everything the API cannot express, and diffs -the result against the registry. Dry run (the default) prints the diff summary and the generated PR body; -``--write`` applies the changes to the root cost map and its ``litellm/`` backup copy. - -Policy highlights: -- Prices arrive per 1M tokens with float artifacts and are normalized to clean per-token values. -- A registry entry absent from the serverless catalog is marked with ``deprecation_date`` from the docs - deprecation table, never deleted; absences with no docs date are surfaced for a human call. -- Availability comes from the API: a model the docs list as removed but the API still serves stays live, - with the conflict surfaced as a warning. -- Manually curated values the API cannot express (``metadata.successor``, ``max_output_tokens`` on existing - entries, capability flags no rule covers) are never overwritten; conflicts are surfaced instead. -""" - -import argparse -import json -import os -import re -import sys -from collections.abc import Mapping, Sequence -from dataclasses import dataclass, field -from pathlib import Path -from types import MappingProxyType -from typing import Final - -import httpx -from pydantic import BaseModel, TypeAdapter, ValidationError - -MODELS_URL: Final = "https://api.together.ai/v1/models?serverless" -DEPRECATIONS_URL: Final = "https://docs.together.ai/docs/deprecations.md" -PROVIDER: Final = "together_ai" -PREFIX: Final = "together_ai/" -SOURCE_URL: Final = "https://docs.together.ai/docs/serverless-models" -COST_MAP_RELPATHS: Final = ( - "model_prices_and_context_window.json", - "litellm/model_prices_and_context_window_backup.json", -) - -TYPE_TO_MODE: Final = MappingProxyType({"chat": "chat", "embedding": "embedding", "moderation": "chat"}) - - -class SyncError(RuntimeError): - pass - - -class CatalogPricing(BaseModel): - input: float - output: float - cached_input: float | None = None - - -class CatalogModel(BaseModel): - id: str - type: str - context_length: int | None = None - pricing: CatalogPricing - - -CATALOG_ADAPTER: Final = TypeAdapter(list[CatalogModel]) - -RegistryEntry = dict[str, object] -CostMap = dict[str, object] - - -@dataclass(frozen=True, slots=True) -class CapabilityRule: - model_id: str - fields: Mapping[str, bool | int] - provenance: str - - -def _rule(model_id: str, provenance: str, **fields: bool | int) -> CapabilityRule: - return CapabilityRule(model_id=model_id, fields=MappingProxyType(dict(fields)), provenance=provenance) - - -_TOOLS: Final = MappingProxyType( - { - "supports_function_calling": True, - "supports_parallel_function_calling": True, - "supports_response_schema": True, - "supports_tool_choice": True, - } -) - -CAPABILITY_RULES: Final = ( - _rule( - "MiniMaxAI/MiniMax-M3", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/minimax-m3", - **_TOOLS, - supports_reasoning=True, - supports_vision=True, - ), - _rule("Prism-ML/Ternary-Bonsai-27B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "Qwen/Qwen3.5-9B", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/qwen3-5-9b", - **_TOOLS, - supports_reasoning=True, - supports_vision=True, - ), - _rule( - "Qwen/Qwen3.6-Plus", - "reviewed for the LIT-5968 backfill; hybrid reasoning model without a documented tools contract", - supports_reasoning=True, - ), - _rule("Qwen/Qwen3.7-Max", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule("Qwen/Qwen3.7-Plus", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule("Qwen/Qwen3.8-2.4T-A95B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule("arize-ai/qwen-2-1.5b-instruct", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "deepseek-ai/DeepSeek-V4-Flash-0731", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-flash", - **_TOOLS, - ), - _rule( - "deepseek-ai/DeepSeek-V4-Pro", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-pro", - **_TOOLS, - supports_reasoning=True, - ), - _rule( - "deepseek-ai/DeepSeek-V4-Pro-0813", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-pro", - **_TOOLS, - ), - _rule("google/gemma-3n-E4B-it", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "google/gemma-4-31B-it", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/gemma-4-31b-it", - **_TOOLS, - supports_vision=True, - ), - _rule( - "intfloat/multilingual-e5-large-instruct", - "embedding dims per https://huggingface.co/intfloat/multilingual-e5-large-instruct", - output_vector_size=1024, - ), - _rule( - "meta-llama/Llama-3.3-70B-Instruct-Turbo", - "reviewed for the LIT-5968 backfill against https://docs.together.ai/docs/function-calling", - **_TOOLS, - ), - _rule( - "meta-llama/Llama-Guard-4-12B", - "moderation classifier with a chat-shaped API; no tools per the LIT-5968 backfill review", - ), - _rule("meta-models/Muse-Glimmer-30B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "moonshotai/Kimi-K2.7-Code", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/kimi-k2-7-code", - **_TOOLS, - supports_vision=True, - ), - _rule( - "moonshotai/Kimi-K3", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/kimi-k3", - **_TOOLS, - supports_reasoning=True, - supports_vision=True, - ), - _rule( - "nvidia/nemotron-3-ultra-550b-a55b", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/nemotron-3-ultra", - **_TOOLS, - supports_reasoning=True, - ), - _rule( - "openai/gpt-oss-120b", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/gpt-oss-120b", - **_TOOLS, - supports_reasoning=True, - ), - _rule( - "openai/gpt-oss-20b", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/gpt-oss-20b", - **_TOOLS, - ), - _rule("pearl-ai/gemma-4-31b-it", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "thinkingmachines/Inkling", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/inkling", - **_TOOLS, - ), - _rule("thinkingmachines/Inkling-Small", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "zai-org/GLM-5.2", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/glm-5-2;" - " 128K output ceiling per https://docs.z.ai/guides/llm/glm-5.2", - **_TOOLS, - supports_reasoning=True, - max_output_tokens=128000, - max_tokens=128000, - ), - _rule( - "zai-org/GLM-5.3-Flash", - "reviewed for LIT-6489 against https://www.together.ai/models/glm-5-3-flash;" - " 128K output ceiling per https://docs.z.ai/guides/llm/glm-5.3", - **_TOOLS, - supports_reasoning=True, - supports_vision=True, - max_output_tokens=128000, - max_tokens=128000, - ), -) - -RULES_BY_ID: Final = MappingProxyType({rule.model_id: rule for rule in CAPABILITY_RULES}) - - -@dataclass(frozen=True, slots=True) -class DeprecationDoc: - removal_dates: Mapping[str, str] - redirects: Mapping[str, str] - - -_REDIRECT_ROW: Final = re.compile(r"^\|\s*`([^`]+)`\s*\|\s*`([^`]+)`\s*\|") -_REMOVAL_ROW: Final = re.compile(r"^\|\s*(\d{4}-\d{2}-\d{2})\s*\|\s*`([^`]+)`\s*\|") - - -def _section(markdown: str, heading: str) -> str: - level: Final = heading.split(" ", 1)[0] - start: Final = markdown.find(f"\n{heading}\n") - if start < 0: - return "" - body: Final = markdown[start + 1 + len(heading) :] - next_heading: Final = re.search(rf"^{re.escape(level)} ", body, flags=re.MULTILINE) - return body[: next_heading.start()] if next_heading else body - - -def parse_deprecations(markdown: str) -> DeprecationDoc: - redirect_rows: Final = tuple( - m.groups() - for m in (_REDIRECT_ROW.match(line) for line in _section(markdown, "## Active model redirects").splitlines()) - if m - ) - inference: Final = _section(_section(markdown, "## Deprecation history"), "### Inference") - removal_rows: Final = tuple(m.groups() for m in (_REMOVAL_ROW.match(line) for line in inference.splitlines()) if m) - if not redirect_rows or not removal_rows: - raise SyncError( - "deprecations doc parsed to zero redirect or removal rows; the table format at " - f"{DEPRECATIONS_URL} changed and the parser needs updating" - ) - removal_dates: Final = {model: date for date, model in reversed(removal_rows)} - return DeprecationDoc( - removal_dates=MappingProxyType(dict(reversed(removal_dates.items()))), - redirects=MappingProxyType({original: target for original, target in redirect_rows}), - ) - - -def per_token(price_per_million: float) -> float: - return float(f"{price_per_million / 1e6:.6g}") - - -def _resolve_name(name: str, universe: frozenset[str]) -> str | None: - if name in universe: - return name - suffix_matches: Final = tuple(candidate for candidate in universe if candidate.endswith(f"/{name}")) - return suffix_matches[0] if len(suffix_matches) == 1 else None - - -def resolve_successor(model_id: str, doc: DeprecationDoc, live_ids: frozenset[str]) -> str | None: - canonical: Final = live_ids | frozenset(doc.removal_dates) - redirects: Final = { - (_resolve_name(raw_source, canonical) or raw_source): (_resolve_name(raw_target, canonical) or raw_target) - for raw_source, raw_target in doc.redirects.items() - } - seen: Final = set() - current = model_id # rebind-ok: walks the redirect chain - while current in redirects and current not in seen: - seen.add(current) - current = redirects[current] # rebind-ok: walks the redirect chain - return current if current != model_id and current in live_ids else None - - -@dataclass(frozen=True, slots=True) -class SyncOutcome: - cost_map: CostMap - added: tuple[str, ...] = () - updated: tuple[str, ...] = () - deprecated: tuple[str, ...] = () - reappeared: tuple[str, ...] = () - warnings: tuple[str, ...] = () - skipped_types: Mapping[str, int] = field(default_factory=dict) - - @property - def has_changes(self) -> bool: - return bool(self.added or self.updated or self.deprecated or self.reappeared) - - -def _api_fields(model: CatalogModel) -> RegistryEntry: - cached: Final = model.pricing.cached_input - return { - "input_cost_per_token": per_token(model.pricing.input), - "output_cost_per_token": per_token(model.pricing.output), - **({"cache_read_input_token_cost": per_token(cached), "supports_prompt_caching": True} if cached else {}), - **({"max_input_tokens": model.context_length} if model.context_length is not None else {}), - } - - -def _new_entry(model: CatalogModel, mode: str) -> RegistryEntry: - rule: Final = RULES_BY_ID.get(model.id) - legacy_ceiling: Final = {} if model.context_length is None else {"max_tokens": model.context_length} - merged: Final = { - **_api_fields(model), - **legacy_ceiling, - "litellm_provider": PROVIDER, - "mode": mode, - "source": SOURCE_URL, - **(dict(rule.fields) if rule else {}), - } - return dict(sorted(merged.items())) - - -def _updated_entry(entry: RegistryEntry, model: CatalogModel) -> tuple[RegistryEntry, tuple[str, ...]]: - rule: Final = RULES_BY_ID.get(model.id) - desired: Final = {**_api_fields(model), **(dict(rule.fields) if rule else {})} - dropped: Final = () if model.pricing.cached_input else ("cache_read_input_token_cost", "supports_prompt_caching") - changes: Final = tuple( - f"{name}: {entry.get(name)!r} -> {value!r}" for name, value in desired.items() if entry.get(name) != value - ) + tuple( - f"{name}: {entry[name]!r} removed (no longer in the catalog pricing)" for name in dropped if name in entry - ) - merged: Final = {name: value for name, value in {**entry, **desired}.items() if name not in dropped} - return dict(sorted(merged.items())), changes - - -def _with_new_keys_in_block(original: CostMap, result: CostMap, new_keys: Sequence[str]) -> CostMap: - provider_keys: Final = tuple(key for key in original if key.startswith(PREFIX)) - if not new_keys or not provider_keys: - return result - block_end: Final = provider_keys[-1] - return { - key: value - for existing in original - for key, value in ( - (existing, result[existing]), - *((new, result[new]) for new in sorted(new_keys) if existing == block_end), - ) - } - - -def compute_sync(cost_map: CostMap, catalog: Sequence[CatalogModel], doc: DeprecationDoc) -> SyncOutcome: - live_ids: Final = frozenset(model.id for model in catalog) - token_models: Final = {model.id: model for model in catalog if model.type in TYPE_TO_MODE} - skipped: Final = { - model.type: sum(1 for m in catalog if m.type == model.type) - for model in catalog - if model.type not in TYPE_TO_MODE - } - registry_ids: Final = {key.removeprefix(PREFIX): key for key in cost_map if key.startswith(PREFIX)} - - added: Final[list[str]] = [] - updated: Final[list[str]] = [] - deprecated: Final[list[str]] = [] - reappeared: Final[list[str]] = [] - warnings: Final[list[str]] = [] - result: Final[CostMap] = dict(cost_map) - - for model_id, model in sorted(token_models.items()): - mode: Final = TYPE_TO_MODE[model.type] - key: Final = f"{PREFIX}{model_id}" - if model_id in doc.removal_dates: - warnings.append( - f"`{key}` is listed as removed on {doc.removal_dates[model_id]} in the docs but the serverless " - "catalog still serves it; availability kept from the API" - ) - entry = result.get(key) - if not isinstance(entry, dict): - result[key] = _new_entry(model, mode) - added.append(key) - if model.type == "chat" and model_id not in RULES_BY_ID: - warnings.append( - f"`{key}` added without a capability rule; review its tools/vision/reasoning support and add one" - ) - continue - if entry.get("mode") != mode: - warnings.append( - f"`{key}` has curated mode {entry.get('mode')!r} but the catalog maps to {mode!r}; left unchanged" - ) - new_entry, changes = _updated_entry(entry, model) - if "deprecation_date" in new_entry: - new_entry.pop("deprecation_date") - reappeared.append(key) - if changes: - updated.append(f"{key}: " + "; ".join(changes)) - if changes or key in reappeared: - result[key] = new_entry - - for model_id, key in sorted(registry_ids.items()): - if model_id in token_models: - continue - entry = result.get(key) - if not isinstance(entry, dict): - continue - removal_date: Final = doc.removal_dates.get(model_id) - successor: Final = resolve_successor(model_id, doc, live_ids) - metadata = entry.get("metadata") - curated_successor: Final = metadata.get("successor") if isinstance(metadata, dict) else None - new_entry = dict(entry) - if removal_date is not None and entry.get("deprecation_date") != removal_date: - if "deprecation_date" in entry: - warnings.append( - f"`{key}` has curated deprecation_date {entry.get('deprecation_date')!r} but the docs list " - f"{removal_date!r}; left unchanged" - ) - else: - new_entry["deprecation_date"] = removal_date - if removal_date is None and "deprecation_date" not in entry: - warnings.append( - f"`{key}` is absent from the serverless catalog with no removal date in the docs; " - "needs a human deprecation call" - ) - if successor is not None: - desired_successor: Final = f"{PREFIX}{successor}" - if curated_successor is None: - new_entry["metadata"] = dict( - sorted({**(metadata if isinstance(metadata, dict) else {}), "successor": desired_successor}.items()) - ) - elif curated_successor != desired_successor: - warnings.append( - f"`{key}` has curated successor {curated_successor!r} but the docs redirects resolve to " - f"{desired_successor!r}; left unchanged" - ) - if new_entry != entry: - result[key] = dict(sorted(new_entry.items())) - deprecated.append(f"{key}: " + ", ".join(sorted(set(new_entry) - set(entry)) or ["updated"])) - - return SyncOutcome( - cost_map=_with_new_keys_in_block(cost_map, result, tuple(added)), - added=tuple(added), - updated=tuple(updated), - deprecated=tuple(deprecated), - reappeared=tuple(reappeared), - warnings=tuple(warnings), - skipped_types=MappingProxyType(skipped), - ) - - -def _section_block(title: str, lines: Sequence[str], backtick: bool) -> str: - bullets: Final = "\n".join(f"- `{line}`" if backtick else f"- {line}" for line in lines) or "- none" - return f"### {title} ({len(lines)})\n{bullets}\n" - - -def render_pr_body(outcome: SyncOutcome) -> str: - skipped: Final = ", ".join(f"{kind} ({count})" for kind, count in sorted(outcome.skipped_types.items())) or "none" - return ( - "Automated daily sync of the together_ai entries in model_prices_and_context_window.json against " - f"`GET {MODELS_URL}` and {DEPRECATIONS_URL} by scripts/sync_together_ai_models.py.\n" - "\n" - f"{_section_block('Added', outcome.added, backtick=True)}" - "\n" - f"{_section_block('Updated', outcome.updated, backtick=True)}" - "\n" - f"{_section_block('Marked deprecated', outcome.deprecated, backtick=True)}" - "\n" - f"{_section_block('Returned to the catalog', outcome.reappeared, backtick=True)}" - "\n" - f"{_section_block('Warnings needing a human call', outcome.warnings, backtick=False)}" - "\n" - f"Catalog model types outside the sync's token-pricing scope, skipped: {skipped}\n" - ) - - -def render_summary(outcome: SyncOutcome) -> str: - return ( - f"added={len(outcome.added)} updated={len(outcome.updated)} deprecated={len(outcome.deprecated)} " - f"reappeared={len(outcome.reappeared)} warnings={len(outcome.warnings)}" - ) - - -def load_catalog(raw: bytes) -> list[CatalogModel]: - parsed: Final = json.loads(raw) - entries: Final = parsed.get("data") if isinstance(parsed, dict) else parsed - try: - catalog: Final = CATALOG_ADAPTER.validate_python(entries) - except ValidationError as error: - raise SyncError(f"the catalog response no longer matches the expected shape: {error}") from error - if not any(model.type in TYPE_TO_MODE for model in catalog): - raise SyncError( - "the catalog response contains no token-priced models; refusing to mark the whole registry deprecated" - ) - return catalog - - -def _fetch(url: str, headers: Mapping[str, str]) -> bytes: - response: Final = httpx.get(url, headers=dict(headers), timeout=30, follow_redirects=True) - if response.status_code != 200: - raise SyncError(f"GET {url} returned {response.status_code}") - return response.content - - -def _serialize(cost_map: CostMap) -> str: - return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n" - - -def main(argv: Sequence[str]) -> int: - parser: Final = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--write", action="store_true", help="apply the sync to the cost map files (default: dry run)") - parser.add_argument("--models-json", type=Path, help="recorded catalog response to use instead of the live API") - parser.add_argument( - "--deprecations-md", type=Path, help="recorded deprecations doc to use instead of the live docs" - ) - parser.add_argument("--pr-body-file", type=Path, help="write the generated PR body to this path") - parser.add_argument("--repo-root", type=Path, default=Path(__file__).resolve().parent.parent) - args: Final = parser.parse_args(argv) - - if args.models_json is not None: - catalog_raw: Final = args.models_json.read_bytes() - else: - api_key: Final = os.environ.get("TOGETHER_API_KEY") - if not api_key: - raise SyncError("TOGETHER_API_KEY is not set and --models-json was not given") - catalog_raw = _fetch(MODELS_URL, {"Authorization": f"Bearer {api_key}"}) # rebind-ok: branch-dependent source - catalog: Final = load_catalog(catalog_raw) - markdown: Final = ( - args.deprecations_md.read_text() if args.deprecations_md is not None else _fetch(DEPRECATIONS_URL, {}).decode() - ) - doc: Final = parse_deprecations(markdown) - - cost_map_path: Final = args.repo_root / COST_MAP_RELPATHS[0] - cost_map: Final = json.loads(cost_map_path.read_text()) - outcome: Final = compute_sync(cost_map, catalog, doc) - body: Final = render_pr_body(outcome) - - if args.pr_body_file is not None: - args.pr_body_file.write_text(body) - if args.write and outcome.has_changes: - for relpath in COST_MAP_RELPATHS: - (args.repo_root / relpath).write_text(_serialize(outcome.cost_map)) - print(render_summary(outcome)) - print() - print(body) - if not args.write: - print("dry run: no files were touched") - elif not outcome.has_changes: - print("registry already in sync: no files were touched") - return 0 - - -if __name__ == "__main__": - try: - raise SystemExit(main(sys.argv[1:])) - except SyncError as error: - print(f"SYNC FAILED: {error}", file=sys.stderr) - raise SystemExit(1) from error diff --git a/tests/unit/test_auto_update_price_and_context_window_file.py b/tests/unit/test_auto_update_price_and_context_window_file.py deleted file mode 100644 index 435747e9a09..00000000000 --- a/tests/unit/test_auto_update_price_and_context_window_file.py +++ /dev/null @@ -1,222 +0,0 @@ -"""Unit tests for the Friendli transform in -`.github/scripts/auto_update_price_and_context_window_file.py`.""" - -from __future__ import annotations - -import importlib.util -import sys -from pathlib import Path - -import pytest - -SCRIPT_PATH = ( - Path(__file__).resolve().parents[2] - / ".github" - / "scripts" - / "auto_update_price_and_context_window_file.py" -) - - -@pytest.fixture(scope="module") -def sync_module(): - spec = importlib.util.spec_from_file_location( - "auto_update_price_and_context_window_file", SCRIPT_PATH - ) - assert spec and spec.loader, f"Could not load spec for {SCRIPT_PATH}" - module = importlib.util.module_from_spec(spec) - sys.modules["auto_update_price_and_context_window_file"] = module - spec.loader.exec_module(module) - return module - - -def _reasoning_model(**overrides: object) -> dict: - model = { - "id": "zai-org/GLM-Test", - "base_model": "zhipuai/glm-test", - "context_length": 1048576, - "max_completion_tokens": 131072, - "pricing": {"input": "0.00000015", "output": "0.0000005", "input_cache_read": "0.00000003"}, - "reasoning": True, - "reasoning_options": [{"type": "effort", "values": ["max", "high", "low"]}], - "functionality": { - "tool_call": True, - "parallel_tool_call": True, - "structured_output": True, - "system_messages": True, - "tool_choice": True, - }, - "input_modalities": ["text", "image", "video"], - "mode": "chat", - } - model.update(overrides) - return model - - -def test_transform_emits_declared_effort_levels_in_canonical_order(sync_module): - entry = sync_module.transform_friendli_data([_reasoning_model()], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert entry["supports_reasoning"] is True - assert entry["reasoning_effort_levels"] == ["low", "high", "max"] - assert not any(k.endswith("_reasoning_effort") for k in entry) - - -def test_transform_reasoning_model_without_effort_options_declares_empty_levels(sync_module): - model = _reasoning_model(reasoning_options=[{"type": "budget_tokens", "values": []}]) - entry = sync_module.transform_friendli_data([model], {})["friendliai/zai-org/GLM-Test"] - assert entry["reasoning_effort_levels"] == [] - - -def test_transform_non_reasoning_model_declares_no_levels(sync_module): - model = _reasoning_model(reasoning=False, reasoning_options=[]) - entry = sync_module.transform_friendli_data([model], {})["friendliai/zai-org/GLM-Test"] - assert entry["supports_reasoning"] is False - assert "reasoning_effort_levels" not in entry - - -def test_transform_max_tokens_mirrors_output_cap_not_context(sync_module): - entry = sync_module.transform_friendli_data([_reasoning_model()], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert entry["max_input_tokens"] == 1048576 - assert entry["max_output_tokens"] == 131072 - assert entry["max_tokens"] == entry["max_output_tokens"] - - -def test_transform_prompt_caching_follows_cache_pricing(sync_module): - cached = sync_module.transform_friendli_data([_reasoning_model()], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert cached["supports_prompt_caching"] is True - assert cached["cache_read_input_token_cost"] == 3e-08 - - uncached_model = _reasoning_model(pricing={"input": "0.00000014", "output": "0.0000004"}) - uncached = sync_module.transform_friendli_data([uncached_model], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert uncached["supports_prompt_caching"] is False - assert "cache_read_input_token_cost" not in uncached - - -def test_transform_modalities_set_vision_image_and_video_flags(sync_module): - entry = sync_module.transform_friendli_data([_reasoning_model()], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert entry["supports_vision"] is True - assert entry["supports_image_input"] is True - assert entry["supports_video_input"] is True - - text_only = _reasoning_model(input_modalities=["text"]) - entry_text = sync_module.transform_friendli_data([text_only], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert entry_text["supports_vision"] is False - assert entry_text["supports_image_input"] is False - assert entry_text["supports_video_input"] is False - - -def test_transform_skips_rows_without_valid_token_prices_so_priced_local_entries_survive(sync_module): - local = { - "friendliai/zai-org/GLM-Test": { - "litellm_provider": "friendliai", - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 5e-07, - } - } - unpriced_rows = [ - _reasoning_model(pricing={}), - _reasoning_model(pricing=None), - _reasoning_model(pricing={"input": "0.00000015"}), - _reasoning_model(pricing={"output": "0.0000005"}), - _reasoning_model(pricing={"input": "not-a-number", "output": "0.0000005"}), - _reasoning_model(pricing={"input": "-0.00000015", "output": "0.0000005"}), - _reasoning_model(pricing={"input": "inf", "output": "0.0000005"}), - _reasoning_model(pricing={"input": "nan", "output": "0.0000005"}), - ] - remote = sync_module.transform_friendli_data(unpriced_rows, local) - assert remote == {} - sync_module.sync_local_data_with_remote(local, remote, replace_keys=frozenset(remote)) - assert local["friendliai/zai-org/GLM-Test"]["input_cost_per_token"] == 1.5e-07 - assert local["friendliai/zai-org/GLM-Test"]["output_cost_per_token"] == 5e-07 - - -def test_transform_keeps_zero_priced_rows(sync_module): - free_model = _reasoning_model(pricing={"input": "0", "output": "0"}) - entry = sync_module.transform_friendli_data([free_model], {})["friendliai/zai-org/GLM-Test"] - assert entry["input_cost_per_token"] == 0.0 - assert entry["output_cost_per_token"] == 0.0 - - -def test_transforms_survive_failed_fetch(sync_module): - assert sync_module.transform_friendli_data(None, {}) == {} - assert sync_module.transform_friendli_data([], {}) == {} - assert sync_module.transform_openrouter_data(None) == {} - assert sync_module.transform_vercel_ai_gateway_data(None) == {} - - -def test_vercel_transform_skips_rows_without_token_pricing_or_limits(sync_module): - rows = [ - { - "id": "wan-video", - "pricing": {"video_duration_pricing": [{"resolution": "720p", "cost_per_second": "0.1"}]}, - }, - { - "id": "qwen3-embedding", - "context_window": 32768, - "max_tokens": 32768, - "pricing": {"input": "0.00000001"}, - }, - { - "id": "no-limits-chat", - "pricing": {"input": "0.000001", "output": "0.000002"}, - }, - { - "id": "good-chat", - "context_window": 128000, - "max_tokens": 8192, - "pricing": {"input": "0.000001", "output": "0.000002"}, - }, - ] - transformed = sync_module.transform_vercel_ai_gateway_data(rows) - assert list(transformed) == ["vercel_ai_gateway/good-chat"] - assert transformed["vercel_ai_gateway/good-chat"]["input_cost_per_token"] == 1e-06 - assert transformed["vercel_ai_gateway/good-chat"]["output_cost_per_token"] == 2e-06 - - -def test_sync_replaces_friendli_entries_so_dropped_cache_pricing_does_not_survive(sync_module): - local = { - "friendliai/zai-org/GLM-Test": { - "litellm_provider": "friendliai", - "cache_read_input_token_cost": 3e-08, - "supports_prompt_caching": True, - } - } - uncached_model = _reasoning_model(pricing={"input": "0.00000014", "output": "0.0000004"}) - remote = sync_module.transform_friendli_data([uncached_model], local) - sync_module.sync_local_data_with_remote(local, remote, replace_keys=frozenset(remote)) - synced = local["friendliai/zai-org/GLM-Test"] - assert "cache_read_input_token_cost" not in synced - assert synced["supports_prompt_caching"] is False - - -def test_sync_still_merges_entries_outside_replace_keys(sync_module): - local = {"openrouter/some-model": {"input_cost_per_token": 1e-06, "supports_vision": True}} - remote = {"openrouter/some-model": {"input_cost_per_token": 2e-06}} - sync_module.sync_local_data_with_remote(local, remote) - assert local["openrouter/some-model"] == {"input_cost_per_token": 2e-06, "supports_vision": True} - - -def test_transform_inherits_allowlisted_keys_from_base_model_entry(sync_module): - local = { - "zhipuai/glm-test": { - "supports_pdf_input": True, - "supports_assistant_prefill": True, - "input_cost_per_token": 9e-06, - } - } - entry = sync_module.transform_friendli_data([_reasoning_model()], local)[ - "friendliai/zai-org/GLM-Test" - ] - assert entry["supports_pdf_input"] is True - assert entry["supports_assistant_prefill"] is True - assert entry["input_cost_per_token"] == 1.5e-07 diff --git a/tests/unit/test_sync_together_ai_models.py b/tests/unit/test_sync_together_ai_models.py deleted file mode 100644 index b8a85bcfbdc..00000000000 --- a/tests/unit/test_sync_together_ai_models.py +++ /dev/null @@ -1,378 +0,0 @@ -import importlib.util -import json -from pathlib import Path -from types import MappingProxyType - -import pytest - -ROOT = Path(__file__).resolve().parents[2] -SCRIPT = ROOT / "scripts" / "sync_together_ai_models.py" -FIXTURES = Path(__file__).resolve().parent / "fixtures" / "together_ai_sync" - -_spec = importlib.util.spec_from_file_location("sync_together_ai_models", SCRIPT) -assert _spec is not None and _spec.loader is not None -sync = importlib.util.module_from_spec(_spec) -_spec.loader.exec_module(sync) - -RECORDED_CATALOG = sync.load_catalog(FIXTURES.joinpath("models_serverless.json").read_bytes()) -RECORDED_DOC = sync.parse_deprecations(FIXTURES.joinpath("deprecations.md").read_text()) - - -def _doc(removal_dates: dict[str, str], redirects: dict[str, str] | None = None) -> object: - return sync.DeprecationDoc( - removal_dates=MappingProxyType(removal_dates), - redirects=MappingProxyType(redirects or {}), - ) - - -def _chat_model(model_id: str, ctx: int = 4096, price: float = 1.0, cached: float | None = None) -> object: - return sync.CatalogModel( - id=model_id, - type="chat", - context_length=ctx, - pricing=sync.CatalogPricing(input=price, output=price, cached_input=cached), - ) - - -@pytest.mark.parametrize( - ("per_million", "expected"), - [ - (3, 3e-06), - (15, 1.5e-05), - (1.4, 1.4e-06), - (0.25999999999999995, 2.6e-07), - (0.060000000000000005, 6e-08), - (1.0399999999999998, 1.04e-06), - (0, 0.0), - ], -) -def test_per_token_normalizes_float_artifacts(per_million: float, expected: float) -> None: - assert sync.per_token(per_million) == expected - - -def test_parse_deprecations_recorded_fixture() -> None: - assert dict(RECORDED_DOC.redirects) == { - "mistralai/Mistral-7B-Instruct-v0.3": "mistralai/Ministral-3-14B-Instruct-2512", - "Kimi-K2": "Kimi-K2-0905", - "DeepSeek-V3": "DeepSeek-V3.1", - "DeepSeek-V3-0324": "DeepSeek-V3.1", - "DeepSeek-R1": "DeepSeek-R1-0528", - } - assert len(RECORDED_DOC.removal_dates) == 208 - assert RECORDED_DOC.removal_dates["google/gemma-3n-E4B-it"] == "2026-08-04" - - -def test_parse_deprecations_duplicate_rows_keep_most_recent_date() -> None: - assert RECORDED_DOC.removal_dates["Qwen/Qwen3-235B-A22B-Thinking-2507"] == "2026-04-16" - - -@pytest.mark.parametrize( - "markdown", - [ - "# Deprecations\n\nNothing here anymore.\n", - "\n## Active model redirects\n\n| A | B |\n| --- | --- |\n| `x` | `y` |\n\n## Something else\n", - "\n## Deprecation history\n\n### Inference\n\n| Date | Model | R |\n| --- | --- | --- |\n| 2026-01-01 | `m` | No |\n", - ], -) -def test_parse_deprecations_raises_when_a_table_parses_empty(markdown: str) -> None: - with pytest.raises(sync.SyncError): - sync.parse_deprecations(markdown) - - -def test_load_catalog_raises_on_shape_change() -> None: - with pytest.raises(sync.SyncError): - sync.load_catalog(b'[{"id": "x", "type": "chat"}]') - - -def test_load_catalog_raises_when_no_token_models_remain() -> None: - only_video = json.dumps([{"id": "v", "type": "video", "pricing": {"input": 0, "output": 0}}]).encode() - with pytest.raises(sync.SyncError): - sync.load_catalog(only_video) - - -def test_recorded_catalog_counts() -> None: - assert len(RECORDED_CATALOG) == 102 - assert sum(1 for model in RECORDED_CATALOG if model.type in sync.TYPE_TO_MODE) == 26 - assert sum(1 for model in RECORDED_CATALOG if model.pricing.cached_input) == 13 - - -def test_added_chat_model_matches_reviewed_registry_shape() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - assert len(outcome.added) == 26 - assert not outcome.deprecated - assert outcome.cost_map["together_ai/moonshotai/Kimi-K3"] == { - "cache_read_input_token_cost": 3e-07, - "input_cost_per_token": 3e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 1048576, - "max_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "source": "https://docs.together.ai/docs/serverless-models", - "supports_function_calling": True, - "supports_parallel_function_calling": True, - "supports_prompt_caching": True, - "supports_reasoning": True, - "supports_response_schema": True, - "supports_tool_choice": True, - "supports_vision": True, - } - - -def test_added_embedding_model_has_no_output_token_cap() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - assert outcome.cost_map["together_ai/intfloat/multilingual-e5-large-instruct"] == { - "input_cost_per_token": 2e-08, - "litellm_provider": "together_ai", - "max_input_tokens": 514, - "max_tokens": 514, - "mode": "embedding", - "output_cost_per_token": 2e-08, - "output_vector_size": 1024, - "source": "https://docs.together.ai/docs/serverless-models", - } - - -def test_moderation_type_maps_to_chat_mode() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - guard = outcome.cost_map["together_ai/meta-llama/Llama-Guard-4-12B"] - assert guard["mode"] == "chat" - assert "max_output_tokens" not in guard - - -def test_output_ceiling_comes_from_the_rule_never_from_context_length() -> None: - glm = next(model for model in RECORDED_CATALOG if model.id == "zai-org/GLM-5.2") - fresh = sync.compute_sync({}, [_chat_model("acme/unreviewed", ctx=1048576), glm], _doc({"x": "2026-01-01"})) - unreviewed = fresh.cost_map["together_ai/acme/unreviewed"] - assert "max_output_tokens" not in unreviewed - assert (unreviewed["max_input_tokens"], unreviewed["max_tokens"]) == (1048576, 1048576) - reviewed = fresh.cost_map["together_ai/zai-org/GLM-5.2"] - assert (reviewed["max_input_tokens"], reviewed["max_output_tokens"], reviewed["max_tokens"]) == ( - 1048575, - 128000, - 128000, - ) - inflated = { - "together_ai/zai-org/GLM-5.2": { - "input_cost_per_token": 1.4e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 1048575, - "max_output_tokens": 1048575, - "max_tokens": 1048575, - "mode": "chat", - "output_cost_per_token": 4.4e-06, - } - } - corrected = sync.compute_sync(inflated, [glm], _doc({"x": "2026-01-01"})) - assert corrected.cost_map["together_ai/zai-org/GLM-5.2"]["max_output_tokens"] == 128000 - assert any("max_output_tokens: 1048575 -> 128000" in line for line in corrected.updated) - - -def test_docs_removed_but_live_model_stays_live_with_warning() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - gemma = outcome.cost_map["together_ai/google/gemma-3n-E4B-it"] - assert "deprecation_date" not in gemma - assert any("gemma-3n-E4B-it" in warning and "2026-08-04" in warning for warning in outcome.warnings) - - -def test_price_change_updates_api_fields_and_keeps_curated_ones() -> None: - registry = { - "together_ai/acme/chat-1": { - "input_cost_per_token": 9e-07, - "litellm_provider": "together_ai", - "max_input_tokens": 4096, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "chat", - "output_cost_per_token": 9e-07, - "supports_audio_input": True, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/chat-1", ctx=8192, price=2.0)], _doc({"x": "2026-01-01"})) - entry = outcome.cost_map["together_ai/acme/chat-1"] - assert entry["input_cost_per_token"] == 2e-06 - assert entry["max_input_tokens"] == 8192 - assert entry["max_output_tokens"] == 2048 - assert entry["supports_audio_input"] is True - assert len(outcome.updated) == 1 - assert "input_cost_per_token" in outcome.updated[0] - - -def test_cached_input_appearing_and_disappearing() -> None: - registry = { - "together_ai/acme/chat-1": { - "cache_read_input_token_cost": 1e-07, - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 1e-06, - "supports_prompt_caching": True, - }, - "together_ai/acme/chat-2": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 1e-06, - }, - } - catalog = [_chat_model("acme/chat-1"), _chat_model("acme/chat-2", cached=0.25999999999999995)] - outcome = sync.compute_sync(registry, catalog, _doc({"x": "2026-01-01"})) - assert "cache_read_input_token_cost" not in outcome.cost_map["together_ai/acme/chat-1"] - assert "supports_prompt_caching" not in outcome.cost_map["together_ai/acme/chat-1"] - assert outcome.cost_map["together_ai/acme/chat-2"]["cache_read_input_token_cost"] == 2.6e-07 - assert outcome.cost_map["together_ai/acme/chat-2"]["supports_prompt_caching"] is True - - -def test_capability_rule_backfills_existing_entry() -> None: - registry = { - "together_ai/moonshotai/Kimi-K3": { - "input_cost_per_token": 3e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - } - } - kimi = next(model for model in RECORDED_CATALOG if model.id == "moonshotai/Kimi-K3") - outcome = sync.compute_sync(registry, [kimi], _doc({"x": "2026-01-01"})) - assert outcome.cost_map["together_ai/moonshotai/Kimi-K3"]["supports_reasoning"] is True - assert any("supports_reasoning" in line for line in outcome.updated) - - -def test_disappeared_model_gets_docs_date_and_is_never_deleted() -> None: - registry = { - "together_ai/acme/gone": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"acme/gone": "2026-07-01"})) - assert outcome.cost_map["together_ai/acme/gone"]["deprecation_date"] == "2026-07-01" - assert outcome.deprecated == ("together_ai/acme/gone: deprecation_date",) - - -def test_disappeared_model_without_docs_date_warns_instead() -> None: - registry = { - "together_ai/acme/gone": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"other": "2026-07-01"})) - assert "deprecation_date" not in outcome.cost_map["together_ai/acme/gone"] - assert not outcome.deprecated - assert any("acme/gone" in warning and "human" in warning for warning in outcome.warnings) - - -def test_curated_deprecation_date_is_never_overwritten() -> None: - registry = { - "together_ai/acme/gone": { - "deprecation_date": "2026-06-15", - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"acme/gone": "2026-07-01"})) - assert outcome.cost_map["together_ai/acme/gone"]["deprecation_date"] == "2026-06-15" - assert any("2026-06-15" in warning and "2026-07-01" in warning for warning in outcome.warnings) - - -def test_redirect_chain_resolves_to_final_live_model() -> None: - doc = _doc({"acme/a": "2026-01-01"}, redirects={"acme/a": "acme/b", "acme/b": "acme/c"}) - live = frozenset({"acme/c"}) - assert sync.resolve_successor("acme/a", doc, live) == "acme/c" - - -def test_redirect_dead_end_yields_no_successor() -> None: - doc = _doc({"acme/a": "2026-01-01"}, redirects={"acme/a": "acme/b"}) - assert sync.resolve_successor("acme/a", doc, frozenset({"acme/other"})) is None - - -def test_redirect_short_names_resolve_by_unique_suffix() -> None: - doc = _doc({"moonshotai/Kimi-K2": "2026-01-01"}, redirects={"Kimi-K2": "Kimi-K2-0905"}) - live = frozenset({"moonshotai/Kimi-K2-0905"}) - assert sync.resolve_successor("moonshotai/Kimi-K2", doc, live) == "moonshotai/Kimi-K2-0905" - - -def test_successor_written_only_when_not_curated() -> None: - registry = { - "together_ai/acme/a": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - }, - "together_ai/acme/b": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "metadata": {"successor": "together_ai/acme/curated"}, - "mode": "chat", - "output_cost_per_token": 1e-06, - }, - } - doc = _doc({"acme/a": "2026-01-01", "acme/b": "2026-01-01"}, redirects={"acme/a": "acme/c", "acme/b": "acme/c"}) - outcome = sync.compute_sync(registry, [_chat_model("acme/c")], doc) - assert outcome.cost_map["together_ai/acme/a"]["metadata"] == {"successor": "together_ai/acme/c"} - assert outcome.cost_map["together_ai/acme/b"]["metadata"] == {"successor": "together_ai/acme/curated"} - assert any("acme/curated" in warning for warning in outcome.warnings) - - -def test_reappearance_clears_deprecation_date() -> None: - registry = { - "together_ai/acme/back": { - "deprecation_date": "2026-05-01", - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 1e-06, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/back")], _doc({"x": "2026-01-01"})) - assert "deprecation_date" not in outcome.cost_map["together_ai/acme/back"] - assert outcome.reappeared == ("together_ai/acme/back",) - - -def test_new_chat_model_without_rule_is_flagged() -> None: - outcome = sync.compute_sync({}, [_chat_model("acme/unreviewed")], _doc({"x": "2026-01-01"})) - assert any("acme/unreviewed" in warning and "capability rule" in warning for warning in outcome.warnings) - - -def test_new_keys_land_at_the_end_of_the_provider_block() -> None: - registry = { - "aaa": {"mode": "chat"}, - "together_ai/acme/old": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - }, - "zzz": {"mode": "chat"}, - } - outcome = sync.compute_sync(registry, [_chat_model("acme/old"), _chat_model("acme/new")], _doc({"x": "2026-01-01"})) - assert list(outcome.cost_map) == ["aaa", "together_ai/acme/old", "together_ai/acme/new", "zzz"] - - -def test_sync_is_idempotent_over_the_repo_cost_map() -> None: - cost_map = json.loads((ROOT / "model_prices_and_context_window.json").read_text()) - first = sync.compute_sync(cost_map, RECORDED_CATALOG, RECORDED_DOC) - second = sync.compute_sync(first.cost_map, RECORDED_CATALOG, RECORDED_DOC) - assert not second.has_changes - assert second.cost_map == first.cost_map - - -def test_pr_body_lists_every_section_and_the_skipped_types() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - body = sync.render_pr_body(outcome) - assert "### Added (26)" in body - assert "### Warnings needing a human call" in body - assert "image (29)" in body - assert "video (38)" in body From b6c3c591eea1f9120d55db4fd6df5a75b3d2f317 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Thu, 29 Oct 2026 07:19:06 -0700 Subject: [PATCH 16/41] ci: run the dashboard API types and Terraform endpoint checks on every change that can alter the spec (#49022) * ci: run the dashboard API types check on every change that can alter the OpenAPI spec * chore: shorten the spec_pattern comment to a pointer at the CI filter * chore: drop the spec_pattern comment * docs: name the full ui-api-types trigger set in the make check header and dashboard AGENTS.md * test: check make check and CI pick the same files for the API-type sync * ci: run the Terraform provider endpoint audit on enterprise route changes * test: keep the Terraform trigger test in the mapped workflow test file --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> (cherry picked from commit e73f06ecdcaf323f21f030c09de42e5a0d6e45e1) --- .github/workflows/test-linting.yml | 4 +- .github/workflows/test-terraform-provider.yml | 1 + scripts/pre_commit_lint.sh | 10 +- tests/unit/test_lint_workflow_diff_gates.py | 100 ++++++++++++++++++ tests/unit/test_pre_commit_lint.py | 28 +++++ ui/litellm-dashboard/AGENTS.md | 2 +- 6 files changed, 136 insertions(+), 9 deletions(-) diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 3b9a2614725..affc4808f8a 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -376,10 +376,10 @@ jobs: exit 0 fi files="$(git diff --name-only "$base" HEAD)" - if grep -Eq '^(litellm/(proxy|types)/|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/test-linting\.yml$)' <<< "$files"; then + if grep -Eq '^(litellm/|enterprise/|(pyproject\.toml|uv\.lock)$|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/test-linting\.yml$)' <<< "$files"; then echo "relevant=true" >> "$GITHUB_OUTPUT" else - echo "No proxy, types or generator changes in this pull request, nothing to verify." + echo "No litellm, enterprise, dependency or generator changes in this pull request, nothing to verify." echo "relevant=false" >> "$GITHUB_OUTPUT" fi diff --git a/.github/workflows/test-terraform-provider.yml b/.github/workflows/test-terraform-provider.yml index ff9db13bd25..6aeb66ed37c 100644 --- a/.github/workflows/test-terraform-provider.yml +++ b/.github/workflows/test-terraform-provider.yml @@ -12,6 +12,7 @@ on: paths: - "terraform/provider/**" - "litellm/proxy/**" + - "enterprise/**" - ".github/workflows/test-terraform-provider.yml" permissions: diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index edb69f24e29..58af67efecc 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -18,7 +18,8 @@ # -> ruff over ruff-tests.toml + `make lint-test-quality` (test-linting.yml's # test-tree ruff and test-quality gate steps) # - dashboard -> prettier + eslint + lint budgets (test-linting.yml's ui job) -# - proxy/types -> regenerate the lazy OpenAPI snapshot and dashboard API types, fail on drift (test-linting.yml's ui-api-types job) +# - spec_pattern -> regenerate the lazy OpenAPI snapshot and dashboard API types, fail on drift +# (test-linting.yml's ui-api-types job) # # Each block is skipped when no matching files are in scope, so unrelated commits # stay fast. This is intentionally not auto-installed as a git hook (see @@ -98,7 +99,7 @@ existing_files() { litellm_py_pattern='^litellm/.*\.py$' e2e_py_pattern='^tests/e2e(_harness)?/.*\.py$' test_tree_pattern='^(tests/.*\.py|ruff-tests\.toml|scripts/(check_test_quality|test_quality_gate)\.py)$' -spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$' +spec_pattern='^(litellm/.*|enterprise/.*|pyproject\.toml|uv\.lock|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$' ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$' ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$' @@ -110,9 +111,6 @@ e2e_py_files=$(scope_match "$e2e_py_pattern") test_tree_files=$(scope_match "$test_tree_pattern") # ruff format (and CI's format step) skip enterprise; the rest of make lint covers it. fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' | existing_files) -# test-linting.yml triggers on any file under litellm/proxy or litellm/types -# (Prisma schema and configs included, not just Python) plus the generator and its -# lockfiles, so match that whole trigger set rather than a Python subset. spec_files=$(scope_match "$spec_pattern") # CI's ui job runs prettier over a wider extension set than eslint; keep that # split so this flags exactly what the job would. @@ -337,7 +335,7 @@ summary_item "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_fi summary_item "test-tree lint (ruff-tests.toml + test-quality gate)" "$test_tree_files" \ "no tests/ Python files or test-tree lint inputs in scope" summary_item "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_changed$ui_eslint_changed" "no dashboard files in scope" -summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/proxy, litellm/types, or generator files in scope" +summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/, enterprise/, dependency, or generator files in scope" if [ -z "$litellm_py_files$e2e_py_files$test_tree_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then echo "check: NOTE - no gating lint check matches the files in scope, so nothing ran:" >&2 diff --git a/tests/unit/test_lint_workflow_diff_gates.py b/tests/unit/test_lint_workflow_diff_gates.py index 7b1d7802d53..5d3c0d80aa6 100644 --- a/tests/unit/test_lint_workflow_diff_gates.py +++ b/tests/unit/test_lint_workflow_diff_gates.py @@ -1,5 +1,6 @@ from __future__ import annotations +import fnmatch import re import shlex import subprocess @@ -7,6 +8,7 @@ from pathlib import Path from typing import Final import pytest +import yaml WORKFLOW: Final = Path(__file__).resolve().parents[2] / ".github" / "workflows" / "test-linting.yml" DIFF_GATE: Final = re.compile(r'git diff --name-only --diff-filter=\w+ "\$GATE_BASE_SHA" HEAD -- (.+?) \|') @@ -100,3 +102,101 @@ def test_e2e_harness_gate_fires_on_any_e2e_or_harness_file_its_dependency_manife (trigger, "tests/e2e/ui/spec.ts", "tests/e2e/ui/pkg/page.py", "elsewhere/pyproject.toml", "tests/e2e_other/module.py"), ) assert selected == frozenset({trigger}) + + +UI_API_TYPES_DETECT: Final[str] = next( + step["run"] + for step in yaml.safe_load(WORKFLOW.read_text())["jobs"]["ui-api-types"]["steps"] + if step.get("id") == "changes" +) +PRE_COMMIT_LINT: Final = WORKFLOW.parents[2] / "scripts" / "pre_commit_lint.sh" +TERRAFORM_PROVIDER_WORKFLOW: Final = WORKFLOW.with_name("test-terraform-provider.yml") + + +def _ui_api_types_relevant(tmp_path: Path, changed: str) -> bool: + _git(tmp_path, "init", "-q", "-b", "main") + _git(tmp_path, "config", "user.email", "t@t") + _git(tmp_path, "config", "user.name", "t") + _git(tmp_path, "commit", "-q", "--allow-empty", "-m", "base") + _git(tmp_path, "switch", "-q", "-c", "pr") + target: Final = tmp_path / changed + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text("x = 1\n") + _git(tmp_path, "add", "-A") + _git(tmp_path, "commit", "-qm", "change") + _git(tmp_path, "switch", "-q", "main") + _git(tmp_path, "commit", "-q", "--allow-empty", "-m", "main moved") + _git(tmp_path, "merge", "-q", "--no-ff", "--no-edit", "pr") + output: Final = tmp_path / "github_output" + subprocess.run( + ["bash", "-c", UI_API_TYPES_DETECT], + cwd=tmp_path, + check=True, + capture_output=True, + env={"PATH": "/usr/bin:/bin", "GITHUB_OUTPUT": str(output)}, + ) + return output.read_text().strip() == "relevant=true" + + +@pytest.mark.parametrize( + "changed", + ( + "litellm/proxy/proxy_server.py", + "litellm/types/router.py", + "litellm/tracing/generated/models.py", + "litellm/router_strategy/complexity_router/config.py", + "litellm/models/credentials.py", + "enterprise/litellm_enterprise/types/proxy/audit_logging_endpoints.py", + "pyproject.toml", + "uv.lock", + "ui/litellm-dashboard/scripts/gen-api-types.mjs", + "ui/litellm-dashboard/package-lock.json", + "ui/litellm-dashboard/src/lib/http/schema.d.ts", + ), +) +def test_ui_api_types_check_runs_on_a_pull_request_touching_any_spec_input(tmp_path: Path, changed: str) -> None: + assert _ui_api_types_relevant(tmp_path, changed) + + +@pytest.mark.parametrize( + "changed", + ( + "litellm-rust/crates/traces/src/lib.rs", + "litellm-proxy-extras/litellm_proxy_extras/utils.py", + "tests/unit/test_router.py", + "ui/litellm-dashboard/src/app/page.tsx", + "docs/my-website/docs/proxy/config.md", + "litellm-dashboard/pyproject.toml", + ), +) +def test_ui_api_types_check_skips_a_pull_request_that_cannot_change_the_spec(tmp_path: Path, changed: str) -> None: + assert not _ui_api_types_relevant(tmp_path, changed) + + +def _first_single_quoted_capture(pattern: str, text: str) -> str: + match: Final = re.search(pattern, text, re.MULTILINE) + assert match is not None, pattern + return match.group(1) + + +def _tracked_paths_matching(extended_regex: str) -> frozenset[str]: + tracked: Final = _git(WORKFLOW.parents[2], "ls-files") + selected: Final = subprocess.run(["grep", "-E", extended_regex], input=tracked, capture_output=True, text=True) + return frozenset(selected.stdout.splitlines()) + + +def test_make_check_runs_the_api_type_sync_on_the_same_tracked_files_as_ci() -> None: + ci_trigger: Final = _first_single_quoted_capture(r"grep -Eq '([^']+)'", UI_API_TYPES_DETECT) + make_check_trigger: Final = _first_single_quoted_capture(r"^spec_pattern='([^']+)'$", PRE_COMMIT_LINT.read_text()) + ci_selected: Final = _tracked_paths_matching(ci_trigger) + assert "pyproject.toml" in ci_selected + assert _tracked_paths_matching(make_check_trigger) == ci_selected - {".github/workflows/test-linting.yml"} + + +def test_a_pull_request_changing_any_route_file_runs_the_terraform_provider_endpoint_audit() -> None: + triggers: Final = yaml.safe_load(TERRAFORM_PROVIDER_WORKFLOW.read_text())[True]["pull_request"]["paths"] + route_files: Final = _git( + WORKFLOW.parents[2], "grep", "-l", "-E", r"APIRouter\(", "--", "*.py", ":(exclude)tests/**" + ).splitlines() + assert route_files + assert [name for name in route_files if not any(fnmatch.fnmatch(name, pattern) for pattern in triggers)] == [] diff --git a/tests/unit/test_pre_commit_lint.py b/tests/unit/test_pre_commit_lint.py index 7c9b98c1c94..ef84ab9fbc8 100644 --- a/tests/unit/test_pre_commit_lint.py +++ b/tests/unit/test_pre_commit_lint.py @@ -427,6 +427,34 @@ def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> No assert "check: FAIL" not in proc.stdout +@pytest.mark.parametrize( + "spec_input", + ( + "litellm/tracing/generated/models.py", + "litellm/models/credentials.py", + "enterprise/litellm_enterprise/types/proxy/audit_logging_endpoints.py", + "pyproject.toml", + "uv.lock", + ), +) +def test_a_staged_spec_input_outside_proxy_and_types_runs_the_api_type_sync(tmp_path: Path, spec_input: str) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _stage_file(repo, spec_input, "x = 1\n") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "ran: dashboard API-type sync (npm run gen:api)" in proc.stdout + + +def test_a_staged_rust_crate_change_skips_the_api_type_sync(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _stage_file(repo, "litellm-rust/crates/traces/src/lib.rs", "fn main() {}\n") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "skipped: dashboard API-type sync (npm run gen:api)" in proc.stdout + + def test_staged_files_matching_no_check_print_an_explicit_noop_note_and_nonempty_log(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) _commit_all(repo, "base") diff --git a/ui/litellm-dashboard/AGENTS.md b/ui/litellm-dashboard/AGENTS.md index b6fe5219022..7e7f48c5462 100644 --- a/ui/litellm-dashboard/AGENTS.md +++ b/ui/litellm-dashboard/AGENTS.md @@ -4,7 +4,7 @@ Never put LiteLLM tokens or API keys in `localStorage`. `localStorage` survives When you fix lint violations that are grandfathered in `eslint-suppressions.json`, run `eslint . --prune-suppressions` and commit the updated baseline so the gate ratchets down instead of leaving a stale suppression -`src/lib/http/schema.d.ts` is generated from the proxy's OpenAPI spec. Never hand-edit it. After changing a backend route or response model that the dashboard consumes, run `npm run gen:api` and commit the result (CI `lint / ui-api-types` enforces this) +`src/lib/http/schema.d.ts` is generated from the proxy's OpenAPI spec. Never hand-edit it. After changing any backend route or request or response model, or bumping a dependency, run `npm run gen:api` and commit the result, since the spec includes every route, not only the ones the dashboard calls (CI `lint / ui-api-types` enforces this) Tests come in three tiers, named by the standard definitions. `Foo.test.tsx` is a unit test: one module, collaborators replaced by doubles, no multi-component tree, and it should run in milliseconds. `Foo.integration.test.tsx` renders a real component tree with real children and only stubs the network boundary. It costs seconds per case, so it earns its place by proving wiring that a unit test cannot reach. Browser-level tests live in `tests/e2e/ui/` as Playwright specs against a live proxy From fea2ca57e3450c2d27b8b43399f0fe199776ff02 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 31 Oct 2026 12:32:23 -0700 Subject: [PATCH 17/41] ci: cut the weekly rc branch on Monday 3am PT instead of Friday (#49472) (cherry picked from commit 3129e4484841a5f714839d444655cfe80ae1e1f9) --- .github/workflows/create-rc-branch.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/create-rc-branch.yml b/.github/workflows/create-rc-branch.yml index 5269460cc93..08ee8351fda 100644 --- a/.github/workflows/create-rc-branch.yml +++ b/.github/workflows/create-rc-branch.yml @@ -2,7 +2,7 @@ name: Create RC Branch on: schedule: - - cron: "0 3 * * 5" + - cron: "0 3 * * 1" timezone: "America/Los_Angeles" workflow_dispatch: From a323682345b461f0bb6b0799a82d64c1f55b8274 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 2 Nov 2026 13:33:05 -0800 Subject: [PATCH 18/41] ci: remove workflows whose secrets are dead (#49722) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit b4a80d175c6f8b1a14bfc22eca934a936a0b8a46) --- .github/workflows/deploy-dev-on-main.yml | 30 -------- .github/workflows/weekly_load_anomaly.yml | 85 ----------------------- tests/e2e/AGENTS.md | 2 +- 3 files changed, 1 insertion(+), 116 deletions(-) delete mode 100644 .github/workflows/deploy-dev-on-main.yml delete mode 100644 .github/workflows/weekly_load_anomaly.yml diff --git a/.github/workflows/deploy-dev-on-main.yml b/.github/workflows/deploy-dev-on-main.yml deleted file mode 100644 index 1208a2dbbdb..00000000000 --- a/.github/workflows/deploy-dev-on-main.yml +++ /dev/null @@ -1,30 +0,0 @@ -name: Deploy main to gateway-dev - -on: - push: - branches: [main] - -permissions: {} - -concurrency: - group: deploy-dev-dispatch - cancel-in-progress: true - -jobs: - dispatch: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - timeout-minutes: 5 - steps: - - uses: actions/create-github-app-token@fee1f7d63c2ff003460e3d139729b119787bc349 # v2.2.2 - id: token - with: - app-id: ${{ secrets.LITELLM_OPS_DISPATCH_APP_ID }} - private-key: ${{ secrets.LITELLM_OPS_DISPATCH_APP_PRIVATE_KEY }} - owner: BerriAI - repositories: litellm-ops - permission-actions: write - - name: Dispatch the litellm-ops dev deploy - env: - GH_TOKEN: ${{ steps.token.outputs.token }} - run: gh api --method POST repos/BerriAI/litellm-ops/actions/workflows/deploy-dev.yml/dispatches -f ref=main diff --git a/.github/workflows/weekly_load_anomaly.yml b/.github/workflows/weekly_load_anomaly.yml deleted file mode 100644 index 0ba4df89cc4..00000000000 --- a/.github/workflows/weekly_load_anomaly.yml +++ /dev/null @@ -1,85 +0,0 @@ -name: "Weekly Load Anomaly Check" - -on: - schedule: - - cron: "0 12 * * 6" - workflow_dispatch: - -permissions: - contents: read - -jobs: - weekly-load-anomaly: - if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - timeout-minutes: 45 - services: - postgres: - image: postgres:16.6 - env: - POSTGRES_USER: llmproxy - POSTGRES_PASSWORD: dbpassword9090 - POSTGRES_DB: litellm - ports: - - 5432:5432 - options: >- - --health-cmd "pg_isready -U llmproxy" - --health-interval 5s - --health-timeout 5s - --health-retries 10 - env: - DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm - LITELLM_MASTER_KEY: sk-weekly-anomaly-check - ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} - AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.AWS_BEARER_TOKEN_BEDROCK }} - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Cache the Rust build - uses: ./.github/actions/cache-cargo-build - - - name: Install dependencies - run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy - - - name: Cache Prisma binaries - uses: ./.github/actions/cache-prisma-binaries - - - name: Generate Prisma client - run: | - uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - - - name: Start the proxy - run: | - nohup uv run --no-sync litellm --config tests/e2e/perf/weekly_anomaly_config.yml --port 4000 > proxy.log 2>&1 & - for _ in $(seq 1 90); do - if curl -fs http://localhost:4000/health/liveliness > /dev/null; then - exit 0 - fi - sleep 2 - done - echo "proxy never became live" - tail -n 100 proxy.log - exit 1 - - - name: Run the weekly session anomaly test - env: - E2E_WEEKLY_ANOMALY: "1" - run: | - uv run --no-sync pytest tests/e2e/perf/test_weekly_session_anomaly_e2e.py -v --tb=short -rA - - - name: Show proxy log on failure - if: failure() - run: tail -n 300 proxy.log diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 88a4df4c40c..ef2ef1c15ad 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -44,7 +44,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) -- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). The memory tests (`test_reliability_memory_e2e.py`, run by the perf lane on its own fresh stack: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches, the check for the v1.100.0 retry-breadcrumb leak). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares) and `Limit`s (one metric at one percentile under one ceiling). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` +- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, run by hand since its weekly workflow never had working Bedrock credentials), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). The memory tests (`test_reliability_memory_e2e.py`, run by the perf lane on its own fresh stack: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches, the check for the v1.100.0 retry-breadcrumb leak). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares) and `Limit`s (one metric at one percentile under one ceiling). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests From 0b2cfa5c1ed57606b57cfb8cd3fcd1454fc7ba50 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 18:16:10 -0800 Subject: [PATCH 19/41] fix(anthropic): keep encrypted reasoning when OpenAI thinking routes the chat adapter to Responses With use_chat_completions_url_for_anthropic_messages on, an openai model with thinking enabled is rewritten to openai/responses/, which accepts the same organization's encrypted reasoning. The keep check now uses the same routing predicate, so turn 2 replays the block instead of stripping it. --- .../pass_through/adapters/handler.py | 34 ++++--------------- .../pass_through/adapters/transformation.py | 30 +++++++++++++++- ...al_pass_through_adapters_transformation.py | 33 ++++++++++++++++++ 3 files changed, 69 insertions(+), 28 deletions(-) diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index e385ee194b3..8b83fe5b7a4 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.tool_fields import ( from litellm.llms.anthropic.common_utils import strip_encrypted_reasoning_blocks_from_anthropic_messages from litellm.llms.anthropic.pass_through.adapters.transformation import ( AnthropicAdapter, + openai_thinking_routes_to_responses, ) from litellm.llms.anthropic.pass_through.context_management import ( AnthropicContextManagementError, @@ -52,7 +53,6 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( ) from litellm.types.llms.openai import OpenAIWebSearchOptions from litellm.types.utils import ModelResponse -from litellm.utils import get_model_info if TYPE_CHECKING: from litellm.chat_completions.dispatch import ChatResult @@ -462,34 +462,14 @@ class LiteLLMMessagesToCompletionTransformationHandler: If the user provides a `summary` field in the thinking dict, it is passed through to the OpenAI reasoning params (opt-in per OpenAI spec). """ - custom_llm_provider = completion_kwargs.get("custom_llm_provider") - if custom_llm_provider is None: - try: - _, inferred_provider, _, _ = litellm.utils.get_llm_provider( - model=cast(str, completion_kwargs.get("model")) - ) - custom_llm_provider = inferred_provider - except Exception: - custom_llm_provider = None - - if custom_llm_provider != "openai": - return - - if not isinstance(thinking, dict) or thinking.get("type") != "enabled": - return - model: Final = completion_kwargs.get("model") - try: - model_info: Final = get_model_info(model=cast(str, model), custom_llm_provider=custom_llm_provider) - if model_info and model_info.get("supports_reasoning") is False: - # Model doesn't support reasoning/responses API, don't route - return - except Exception: - pass + if not isinstance(model, str) or thinking is None: + return + if not openai_thinking_routes_to_responses(model, completion_kwargs.get("custom_llm_provider"), thinking): + return - if isinstance(model, str) and model and "responses/" not in model: - local_model: Final = model.removeprefix(f"{custom_llm_provider}/") - completion_kwargs["model"] = f"{custom_llm_provider}/responses/{local_model}" + if "responses/" not in model: + completion_kwargs["model"] = f"openai/responses/{model.removeprefix('openai/')}" auto_summary: Final = is_reasoning_auto_summary_enabled() diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index ae7c433de2a..ee553749c74 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -227,11 +227,36 @@ class ChatRouteInputs: tools: Sequence[Mapping[str, object]] | None = None reasoning_effort: str | Mapping[str, object] | None = None api_base: str | None = None + thinking: Mapping[str, object] | None = None _NO_ROUTE_INPUTS: Final = ChatRouteInputs() +def _inferred_provider(model: str) -> str | None: + try: + return litellm.utils.get_llm_provider(model=model)[1] + except Exception: # noqa: BLE001 # an unknown model name has no inferred provider + return None + + +def openai_thinking_routes_to_responses( + model: str | None, + custom_llm_provider: str | None, + thinking: Mapping[str, object] | None, +) -> bool: + if not model or thinking is None or thinking.get("type") != "enabled": + return False + provider: Final = custom_llm_provider or _inferred_provider(model) + if provider != "openai": + return False + try: + model_info: Final = litellm.get_model_info(model=model, custom_llm_provider=provider) + except Exception: # noqa: BLE001 # a model missing from the cost map still routes, as before + return True + return model_info.get("supports_reasoning") is not False + + def _keeps_encrypted_reasoning_for_responses( model: str | None, custom_llm_provider: str | None, @@ -239,7 +264,9 @@ def _keeps_encrypted_reasoning_for_responses( ) -> bool: from litellm.responses.bridge_check import chat_completion_ends_in_responses_api - return chat_completion_ends_in_responses_api( + return openai_thinking_routes_to_responses( + model, custom_llm_provider, route.thinking + ) or chat_completion_ends_in_responses_api( model, custom_llm_provider, tools=route.tools, @@ -1321,6 +1348,7 @@ class LiteLLMAnthropicMessagesAdapter: tools=new_kwargs.get("tools"), reasoning_effort=new_kwargs.get("reasoning_effort"), api_base=api_base, + thinking=anthropic_message_request.get("thinking"), ), ) ## ADD SYSTEM MESSAGE TO MESSAGES diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index bc357878944..7a8caa417df 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -5566,3 +5566,36 @@ async def test_messages_do_not_retry_other_rejections( ) assert route.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_openai_thinking_on_the_chat_adapter_keeps_reasoning_its_responses_route_replays( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "use_chat_completions_url_for_anthropic_messages", True) + route: Final = respx_mock.post(url__regex=r".*openai-thinking\.example/v1/responses").mock( + return_value=httpx.Response(200, json=_RETRY_REPLY) + ) + + await litellm.anthropic.messages.acreate( + model="openai/gpt-5-mini", + api_base="https://openai-thinking.example/v1", + api_key="fake-key", + max_tokens=2048, + thinking={"type": "enabled", "budget_tokens": 1024}, + messages=[ + {"role": "user", "content": "Which city is colder?"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAA_same_org")}, + {"type": "text", "text": "Oslo."}, + ], + }, + {"role": "user", "content": "And by how much?"}, + ], + ) + + assert route.call_count == 1 + assert [item["encrypted_content"] for item in _sent_reasoning_items(route.calls[0].request)] == ["gAAA_same_org"] From 4fc1d5bb0d077898f46b2649abe1c805abbf8354 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 18:16:29 -0800 Subject: [PATCH 20/41] fix(anthropic): stop logging rejected encrypted reasoning The retry warning keeps provider, model, status, and error code, and leaves the ciphertext out. Backport of the third #49076 commit to lts/1.108.x. Cherry-picked from c62c92d2ad0. Co-authored-by: DDT <141296842+Hugo-DDT@users.noreply.github.com> --- .../pass_through/adapters/handler.py | 37 +++++++++++++++++-- ...al_pass_through_adapters_transformation.py | 34 ++++++++++++----- 2 files changed, 59 insertions(+), 12 deletions(-) diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 8b83fe5b7a4..963dd050804 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -83,15 +83,46 @@ class _ChatTargetReasoning: summary: str | None = None +def _provider_error_code(body: object) -> str | None: + mapping: Final = as_str_mapping(body) + if mapping is None: + return None + nested: Final = as_str_mapping(mapping.get("error")) + source: Final = nested if nested is not None else mapping + code: Final = source.get("code") + return code if isinstance(code, str) else None + + +def _logged_rejection_code(error: litellm.BadRequestError) -> str: + if error.code == "invalid_encrypted_content" or _provider_error_code(error.body) == "invalid_encrypted_content": + return "invalid_encrypted_content" + if "invalid_encrypted_content" in str(error): + return "invalid_encrypted_content" + return "could_not_be_verified" + + +def _completion_text(completion_kwargs: _CompletionKwargs, key: str) -> str: + value: Final = completion_kwargs.get(key) + return value if isinstance(value, str) else "" + + def _messages_without_rejected_reasoning( messages: _AnthropicMessages, error: litellm.BadRequestError, + completion_kwargs: _CompletionKwargs, ) -> "_AnthropicMessages | None": error_message: Final = str(error) rejected: Final = "invalid_encrypted_content" in error_message or "could not be verified" in error_message if not rejected or not messages_carry_encrypted_reasoning(messages): return None - verbose_logger.warning("Backend rejected replayed encrypted reasoning, retrying once without it: %s", error) + verbose_logger.warning( + "Backend rejected replayed encrypted reasoning, retrying once without it. " + "provider=%s model=%s status=%s code=%s", + _completion_text(completion_kwargs, "custom_llm_provider"), + _completion_text(completion_kwargs, "model"), + error.status_code, + _logged_rejection_code(error), + ) return _ANTHROPIC_MESSAGES.validate_python(strip_encrypted_reasoning_blocks_from_anthropic_messages(messages)) @@ -103,7 +134,7 @@ async def _acompletion_without_rejected_reasoning( try: return await litellm.acompletion(**completion_kwargs) except litellm.BadRequestError as error: - replayed_messages: Final = _messages_without_rejected_reasoning(messages, error) + replayed_messages: Final = _messages_without_rejected_reasoning(messages, error, completion_kwargs) if replayed_messages is None: raise return await litellm.acompletion(**prepare_replay(replayed_messages)) @@ -117,7 +148,7 @@ def _completion_without_rejected_reasoning( try: return litellm.completion(**completion_kwargs) except litellm.BadRequestError as error: - replayed_messages: Final = _messages_without_rejected_reasoning(messages, error) + replayed_messages: Final = _messages_without_rejected_reasoning(messages, error, completion_kwargs) if replayed_messages is None: raise return litellm.completion(**prepare_replay(replayed_messages)) diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 7a8caa417df..3d87772595a 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1,6 +1,7 @@ import base64 import functools import json +import logging from types import MappingProxyType from typing import Any, Final, cast @@ -5492,25 +5493,40 @@ def _sent_reasoning_items(request: httpx.Request) -> list[object]: @pytest.mark.usefixtures("_chat_bridge_httpx") async def test_async_messages_retry_once_without_reasoning_the_backend_cannot_decrypt( respx_mock: respx.MockRouter, + caplog: pytest.LogCaptureFixture, ) -> None: route: Final = respx_mock.post(url__regex=_RETRY_URL).mock( side_effect=[_RETRY_REJECTION, httpx.Response(200, json=_RETRY_REPLY)] ) - response: Final = await litellm.anthropic.messages.acreate( - model="azure/gpt-5.4-nano", - api_base=_RETRY_API_BASE, - api_key="fake-key", - api_version="2025-04-01-preview", - max_tokens=64, - tools=_RETRY_TOOLS, - messages=_history_with_foreign_reasoning(), - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + response: Final = await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + ) + warnings: Final = tuple( + record.getMessage() + for record in caplog.records + if record.levelno >= logging.WARNING and record.name == "LiteLLM" + ) + retry_warnings: Final = tuple(message for message in warnings if "retrying once without it" in message) assert route.call_count == 2 assert len(_sent_reasoning_items(route.calls[0].request)) == 1 assert _sent_reasoning_items(route.calls[1].request) == [] assert [block["text"] for block in response["content"] if block["type"] == "text"] == ["Oslo"] + assert retry_warnings == ( + "Backend rejected replayed encrypted reasoning, retrying once without it. " + "provider=azure model=azure/gpt-5.4-nano status=400 code=invalid_encrypted_content", + ) + assert all( + "gAAA_foreign" not in message and "litellm_encrypted_reasoning:" not in message for message in warnings + ) @pytest.mark.usefixtures("_chat_bridge_httpx") From ba5d7b83b950b9913847e0a63e0cc7afdf5fbeb8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 19:30:43 -0800 Subject: [PATCH 21/41] fix(anthropic): start a new streamed block for each reasoning item The backported Responses bridge tags each reasoning item's thinking block with its output index, but this line's Anthropic stream wrapper only opened a new block when the block type changed, so a second reasoning item of the same type was dropped. This ports the index check from #46863, which main already has. --- .../adapters/streaming_iterator.py | 23 ++++++ ...al_pass_through_adapters_transformation.py | 73 +++++++++++++++++++ 2 files changed, 96 insertions(+) diff --git a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py index cc705ab1184..d8f4ad0b802 100644 --- a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py @@ -24,6 +24,7 @@ from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.exceptions import ContextWindowExceededError, MidStreamFallbackError +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.hidden_params import set_hidden_params from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame @@ -71,6 +72,14 @@ def _optional_attr_sequence(obj: object, name: str) -> Sequence[object]: return value if value else () +def _first_thinking_block_index(chunk: "ModelResponseStream") -> int | None: + delta: Final = _optional_attr(chunk.choices[0], "delta") if chunk.choices else None + blocks: Final = _optional_attr_sequence(delta, "thinking_blocks") + first_block: Final = as_str_mapping(blocks[0]) if blocks else None + index: Final = first_block.get("index") if first_block is not None else None + return index if isinstance(index, int) else None + + def _error_status_and_message(exc: Exception) -> tuple[int, str]: if isinstance(exc, (BaseLLMException, MidStreamFallbackError, ContextWindowExceededError)): return exc.status_code, exc.message @@ -464,6 +473,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): holding_stop_reason_chunk: MessageBlockDelta | None = None queued_usage_chunk: bool = False current_content_block_index: int = 0 + current_thinking_block_index: int | None = None def __init__( self, @@ -1352,11 +1362,24 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): original_name: Final = self.tool_name_mapping.get(truncated_name, truncated_name) tool_block["name"] = original_name + thinking_block_index: Final = _first_thinking_block_index(chunk) + previous_thinking_block_index: Final = self.current_thinking_block_index + if thinking_block_index is not None or block_type != self.current_content_block_type: + self.current_thinking_block_index = thinking_block_index + if block_type != self.current_content_block_type: self.current_content_block_type = block_type self.current_content_block_start = content_block_start return True + if ( + thinking_block_index is not None + and previous_thinking_block_index is not None + and thinking_block_index != previous_thinking_block_index + ): + self.current_content_block_start = content_block_start + return True + # For parallel tool calls, we'll necessarily have a new content block # if we get a function name since it signals a new tool call if block_type == "tool_use": diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 3d87772595a..6411d57d1bd 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -2,6 +2,7 @@ import base64 import functools import json import logging +from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import Any, Final, cast @@ -5615,3 +5616,75 @@ async def test_openai_thinking_on_the_chat_adapter_keeps_reasoning_its_responses assert route.call_count == 1 assert [item["encrypted_content"] for item in _sent_reasoning_items(route.calls[0].request)] == ["gAAA_same_org"] + + +def _responses_sse(events: Sequence[Mapping[str, object]]) -> bytes: + return b"".join( + f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': number})}\n\n".encode() + for number, event in enumerate(events) + ) + + +def _reasoning_item(item_id: str, encrypted: str) -> dict[str, object]: + return {"type": "reasoning", "id": item_id, "summary": [], "encrypted_content": encrypted} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_streamed_consecutive_reasoning_items_each_become_their_own_block( + respx_mock: respx.MockRouter, +) -> None: + response: Final = {**_RETRY_REPLY, "id": "resp_two_reasoning", "status": "in_progress", "output": []} + message: Final = { + "type": "message", + "id": "msg_two_reasoning", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Oslo", "annotations": []}], + } + first: Final = _reasoning_item("rs_1", "gAAA_first") + second: Final = _reasoning_item("rs_2", "gAAA_second") + completed: Final = {**response, "status": "completed", "output": [first, second, message]} + events: Final = ( + {"type": "response.created", "response": response}, + {"type": "response.output_item.added", "output_index": 0, "item": _reasoning_item("rs_1", "")}, + {"type": "response.output_item.done", "output_index": 0, "item": first}, + {"type": "response.output_item.added", "output_index": 1, "item": _reasoning_item("rs_2", "")}, + {"type": "response.output_item.done", "output_index": 1, "item": second}, + {"type": "response.output_item.added", "output_index": 2, "item": {**message, "content": []}}, + { + "type": "response.output_text.delta", + "output_index": 2, + "item_id": "msg_two_reasoning", + "content_index": 0, + "delta": "Oslo", + }, + {"type": "response.output_item.done", "output_index": 2, "item": message}, + {"type": "response.completed", "response": completed}, + ) + respx_mock.post(url__regex=_RETRY_URL).mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=_responses_sse(events)) + ) + + stream: Final = await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + stream=True, + messages=[{"role": "user", "content": "Which city is colder?"}], + ) + frames: Final = b"".join([chunk async for chunk in stream]).decode() + starts: Final = tuple( + json.loads(line.removeprefix("data: "))["content_block"] + for line in frames.splitlines() + if line.startswith("data: ") and '"content_block_start"' in line + ) + + assert [(block["type"], block.get("data")) for block in starts] == [ + ("redacted_thinking", encrypted_reasoning_signature("gAAA_first")), + ("redacted_thinking", encrypted_reasoning_signature("gAAA_second")), + ("text", None), + ] From 27574a429a54903c8708f3153e06a667b7bf69cf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 20:45:06 -0800 Subject: [PATCH 22/41] fix(anthropic): keep each streamed reasoning summary with its own item A summary delta carries no item index, so the next item's summary joined the block that already got its signature. A thinking chunk after an indexed signature now opens a new block. --- .../adapters/streaming_iterator.py | 27 +++-- ...al_pass_through_adapters_transformation.py | 100 +++++++++++++++++- 2 files changed, 119 insertions(+), 8 deletions(-) diff --git a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py index d8f4ad0b802..6f094256197 100644 --- a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py @@ -72,14 +72,21 @@ def _optional_attr_sequence(obj: object, name: str) -> Sequence[object]: return value if value else () -def _first_thinking_block_index(chunk: "ModelResponseStream") -> int | None: +def _first_thinking_block(chunk: "ModelResponseStream") -> Mapping[str, object] | None: delta: Final = _optional_attr(chunk.choices[0], "delta") if chunk.choices else None blocks: Final = _optional_attr_sequence(delta, "thinking_blocks") - first_block: Final = as_str_mapping(blocks[0]) if blocks else None - index: Final = first_block.get("index") if first_block is not None else None + return as_str_mapping(blocks[0]) if blocks else None + + +def _thinking_block_index(block: Mapping[str, object] | None) -> int | None: + index: Final = block.get("index") if block is not None else None return index if isinstance(index, int) else None +def _closes_thinking_block(block: Mapping[str, object] | None) -> bool: + return block is not None and _thinking_block_index(block) is not None and bool(block.get("signature")) + + def _error_status_and_message(exc: Exception) -> tuple[int, str]: if isinstance(exc, (BaseLLMException, MidStreamFallbackError, ContextWindowExceededError)): return exc.status_code, exc.message @@ -474,6 +481,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): queued_usage_chunk: bool = False current_content_block_index: int = 0 current_thinking_block_index: int | None = None + current_thinking_block_closed: bool = False def __init__( self, @@ -1362,9 +1370,16 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): original_name: Final = self.tool_name_mapping.get(truncated_name, truncated_name) tool_block["name"] = original_name - thinking_block_index: Final = _first_thinking_block_index(chunk) + thinking_block: Final = _first_thinking_block(chunk) + thinking_block_index: Final = _thinking_block_index(thinking_block) previous_thinking_block_index: Final = self.current_thinking_block_index - if thinking_block_index is not None or block_type != self.current_content_block_type: + follows_closed_thinking_block: Final = self.current_thinking_block_closed and block_type == "thinking" + self.current_thinking_block_closed = _closes_thinking_block(thinking_block) + if ( + thinking_block_index is not None + or block_type != self.current_content_block_type + or follows_closed_thinking_block + ): self.current_thinking_block_index = thinking_block_index if block_type != self.current_content_block_type: @@ -1372,7 +1387,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): self.current_content_block_start = content_block_start return True - if ( + if follows_closed_thinking_block or ( thinking_block_index is not None and previous_thinking_block_index is not None and thinking_block_index != previous_thinking_block_index diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 6411d57d1bd..cda53b82931 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -5625,8 +5625,13 @@ def _responses_sse(events: Sequence[Mapping[str, object]]) -> bytes: ) -def _reasoning_item(item_id: str, encrypted: str) -> dict[str, object]: - return {"type": "reasoning", "id": item_id, "summary": [], "encrypted_content": encrypted} +def _reasoning_item(item_id: str, encrypted: str, summary: str = "") -> dict[str, object]: + return { + "type": "reasoning", + "id": item_id, + "summary": [{"type": "summary_text", "text": summary}] if summary else [], + "encrypted_content": encrypted, + } @pytest.mark.asyncio @@ -5688,3 +5693,94 @@ async def test_streamed_consecutive_reasoning_items_each_become_their_own_block( ("redacted_thinking", encrypted_reasoning_signature("gAAA_second")), ("text", None), ] + + +def _streamed_blocks(frames: str) -> list[tuple[str, str, str | None]]: + events: Final = tuple( + json.loads(line.removeprefix("data: ")) for line in frames.splitlines() if line.startswith("data: ") + ) + deltas: Final = tuple(event for event in events if event["type"] == "content_block_delta") + + def block_deltas(index: int) -> tuple[Mapping[str, str], ...]: + return tuple(delta["delta"] for delta in deltas if delta["index"] == index) + + return [ + ( + event["content_block"]["type"], + "".join(delta.get("thinking", "") for delta in block_deltas(event["index"])), + next( + (delta["signature"] for delta in block_deltas(event["index"]) if delta["type"] == "signature_delta"), + None, + ), + ) + for event in events + if event["type"] == "content_block_start" + ] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_streamed_reasoning_summaries_stay_with_their_own_encrypted_item( + respx_mock: respx.MockRouter, +) -> None: + response: Final = {**_RETRY_REPLY, "id": "resp_two_summaries", "status": "in_progress", "output": []} + message: Final = { + "type": "message", + "id": "msg_two_summaries", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Oslo", "annotations": []}], + } + first: Final = _reasoning_item("rs_1", "gAAA_first", "Oslo is further north.") + second: Final = _reasoning_item("rs_2", "gAAA_second", "Check the forecast.") + completed: Final = {**response, "status": "completed", "output": [first, second, message]} + + def summary_delta(output_index: int, item_id: str, text: str) -> dict[str, object]: + return { + "type": "response.reasoning_summary_text.delta", + "output_index": output_index, + "item_id": item_id, + "summary_index": 0, + "delta": text, + } + + events: Final = ( + {"type": "response.created", "response": response}, + {"type": "response.output_item.added", "output_index": 0, "item": _reasoning_item("rs_1", "")}, + summary_delta(0, "rs_1", "Oslo is further north."), + {"type": "response.output_item.done", "output_index": 0, "item": first}, + {"type": "response.output_item.added", "output_index": 1, "item": _reasoning_item("rs_2", "")}, + summary_delta(1, "rs_2", "Check the forecast."), + {"type": "response.output_item.done", "output_index": 1, "item": second}, + {"type": "response.output_item.added", "output_index": 2, "item": {**message, "content": []}}, + { + "type": "response.output_text.delta", + "output_index": 2, + "item_id": "msg_two_summaries", + "content_index": 0, + "delta": "Oslo", + }, + {"type": "response.output_item.done", "output_index": 2, "item": message}, + {"type": "response.completed", "response": completed}, + ) + respx_mock.post(url__regex=_RETRY_URL).mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=_responses_sse(events)) + ) + + stream: Final = await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + stream=True, + messages=[{"role": "user", "content": "Which city is colder?"}], + ) + frames: Final = b"".join([chunk async for chunk in stream]).decode() + + assert _streamed_blocks(frames) == [ + ("thinking", "Oslo is further north.", encrypted_reasoning_signature("gAAA_first")), + ("thinking", "Check the forecast.", encrypted_reasoning_signature("gAAA_second")), + ("text", "", None), + ] From f86222be48c20554817f8a35e1e83cd3216a0854 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 23:23:15 -0800 Subject: [PATCH 23/41] fix(anthropic): stream reasoning signatures only to Messages clients, keep xAI web search reasoning, and retry only when the sent request carried it --- .../handler.py | 10 ++ .../transformation.py | 15 ++- .../pass_through/adapters/handler.py | 25 ++++- .../pass_through/adapters/transformation.py | 4 + litellm/responses/bridge_check.py | 2 + ...responses_transformation_transformation.py | 104 +++++++++++++++++- ...al_pass_through_adapters_transformation.py | 56 ++++++++++ 7 files changed, 206 insertions(+), 10 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index 3f7017eccc2..eef9cc6c4c9 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -13,7 +13,9 @@ from typing import ( from typing_extensions import TypedDict +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import CallTypes if TYPE_CHECKING: from litellm import CustomStreamWrapper, LiteLLMLoggingObj, ModelResponse @@ -36,6 +38,10 @@ def _restore_routing_prefix(model: str, custom_llm_provider: str) -> str: return f"{custom_llm_provider}/{model}" +def _serves_anthropic_messages(logging_obj: "LiteLLMLoggingObj") -> bool: + return attribute_or_none(logging_obj, "call_type") == CallTypes.anthropic_messages.value + + class ResponsesToCompletionBridgeHandler: def __init__(self): from .transformation import LiteLLMResponsesTransformationHandler @@ -182,6 +188,7 @@ class ResponsesToCompletionBridgeHandler: model_response: Final = validated_kwargs["model_response"] logging_obj: Final = validated_kwargs["logging_obj"] custom_llm_provider: Final = validated_kwargs["custom_llm_provider"] + serves_anthropic_messages: Final = _serves_anthropic_messages(logging_obj) if kwargs.get("stream") is True and "stream" not in optional_params: optional_params = {**optional_params, "stream": True} @@ -260,6 +267,7 @@ class ResponsesToCompletionBridgeHandler: structured_output=is_structured_output_request( cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict ), + streams_reasoning_signatures=serves_anthropic_messages, ) streamwrapper: Final = CustomStreamWrapper( completion_stream=completion_stream, @@ -282,6 +290,7 @@ class ResponsesToCompletionBridgeHandler: model_response: Final = validated_kwargs["model_response"] logging_obj: Final = validated_kwargs["logging_obj"] custom_llm_provider: Final = validated_kwargs["custom_llm_provider"] + serves_anthropic_messages: Final = _serves_anthropic_messages(logging_obj) if kwargs.get("stream") is True and "stream" not in optional_params: optional_params = {**optional_params, "stream": True} @@ -363,6 +372,7 @@ class ResponsesToCompletionBridgeHandler: structured_output=is_structured_output_request( cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict ), + streams_reasoning_signatures=serves_anthropic_messages, ) streamwrapper: Final = CustomStreamWrapper( completion_stream=completion_stream, diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 296f1949d11..33ac222b06b 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -1563,9 +1563,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): sync_stream: bool, json_mode: bool | None = False, structured_output: bool = False, + streams_reasoning_signatures: bool = False, ) -> BaseModelResponseIterator: return OpenAiResponsesToChatCompletionStreamIterator( - streaming_response, sync_stream, json_mode, structured_output=structured_output + streaming_response, + sync_stream, + json_mode, + structured_output=structured_output, + streams_reasoning_signatures=streams_reasoning_signatures, ) def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, object]: @@ -1944,8 +1949,10 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): sync_stream: bool, json_mode: bool | None = False, structured_output: bool = False, + streams_reasoning_signatures: bool = False, ): super().__init__(streaming_response, sync_stream, json_mode) + self._streams_reasoning_signatures: Final = streams_reasoning_signatures self._chat_completion_id: str | None = None self._served_service_tier: str | None = None self._tool_call_index_map: dict[int, int] = {} # mutable-ok: per-stream accumulator state @@ -2266,7 +2273,11 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): """ verbose_logger.debug("Chat provider: transform_streaming_response called with chunk: %s", chunk) self._remember_served_service_tier(chunk) - reasoning_chunk: Final = _reasoning_done_stream_chunk(chunk, self._summary_streamed_outputs) + reasoning_chunk: Final = ( + _reasoning_done_stream_chunk(chunk, self._summary_streamed_outputs) + if self._streams_reasoning_signatures + else None + ) translated: Final = ( self._empty_content_chunk() if self._is_text_delta_after_structured_answer(chunk) diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 963dd050804..739077d3566 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -14,9 +14,12 @@ from typing_extensions import TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import run_async_function -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject -from litellm.litellm_core_utils.prompt_templates.common_utils import messages_carry_encrypted_reasoning +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + is_encrypted_reasoning_block, + thinking_blocks_with_encrypted_reasoning, +) from litellm.litellm_core_utils.tool_fields import ( ROUTES_WITHOUT_MCP_SERVERS, drops_mcp_servers, @@ -106,6 +109,22 @@ def _completion_text(completion_kwargs: _CompletionKwargs, key: str) -> str: return value if isinstance(value, str) else "" +def _message_carries_encrypted_reasoning(message: object) -> bool: + mapping: Final = as_str_mapping(message) + if mapping is None: + return False + blocks: Final = thinking_blocks_with_encrypted_reasoning( + mapping.get("thinking_blocks"), mapping.get("reasoning_items") + ) + return any(is_encrypted_reasoning_block(block) for block in blocks) + + +def _sends_encrypted_reasoning(completion_kwargs: _CompletionKwargs) -> bool: + messages: Final = completion_kwargs.get("messages") + sent: Final = iter_items(messages) if messages is not None else () + return any(_message_carries_encrypted_reasoning(message) for message in sent) + + def _messages_without_rejected_reasoning( messages: _AnthropicMessages, error: litellm.BadRequestError, @@ -113,7 +132,7 @@ def _messages_without_rejected_reasoning( ) -> "_AnthropicMessages | None": error_message: Final = str(error) rejected: Final = "invalid_encrypted_content" in error_message or "could not be verified" in error_message - if not rejected or not messages_carry_encrypted_reasoning(messages): + if not rejected or not _sends_encrypted_reasoning(completion_kwargs): return None verbose_logger.warning( "Backend rejected replayed encrypted reasoning, retrying once without it. " diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index ee553749c74..0f0f7cd22f0 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -201,6 +201,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolParamFunctionChunk, ChatCompletionToolReferenceObject, ChatCompletionUserMessage, + OpenAIWebSearchOptions, ToolMessageContentPart, ) from litellm.types.utils import Choices, ModelResponse, StreamingChoices, Usage @@ -228,6 +229,7 @@ class ChatRouteInputs: reasoning_effort: str | Mapping[str, object] | None = None api_base: str | None = None thinking: Mapping[str, object] | None = None + web_search: bool = False _NO_ROUTE_INPUTS: Final = ChatRouteInputs() @@ -272,6 +274,7 @@ def _keeps_encrypted_reasoning_for_responses( tools=route.tools, reasoning_effort=route.reasoning_effort, api_base=route.api_base, + web_search_options=OpenAIWebSearchOptions() if route.web_search else None, ) @@ -1349,6 +1352,7 @@ class LiteLLMAnthropicMessagesAdapter: reasoning_effort=new_kwargs.get("reasoning_effort"), api_base=api_base, thinking=anthropic_message_request.get("thinking"), + web_search="web_search_options" in new_kwargs, ), ) ## ADD SYSTEM MESSAGE TO MESSAGES diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index d8a9440bc57..4a25bc815ba 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -198,6 +198,7 @@ def chat_completion_ends_in_responses_api( tools: Sequence[Mapping[str, object]] | None = None, reasoning_effort: str | Mapping[str, object] | None = None, api_base: str | None = None, + web_search_options: OpenAIWebSearchOptions | None = None, ) -> bool: """Whether chat completion for this model is served by a Responses API. @@ -217,6 +218,7 @@ def chat_completion_ends_in_responses_api( model_info, _resolved = responses_api_bridge_check( model=resolved_model, custom_llm_provider=provider, + web_search_options=web_search_options, tools=tools, reasoning_effort=reasoning_effort, api_base=api_base, diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index d1155e41bdf..df9d048fb45 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3267,7 +3267,7 @@ def test_reasoning_output_item_done_emits_signed_thinking_blocks_the_stream_keep }, } result: Final = OpenAiResponsesToChatCompletionStreamIterator( - streaming_response=None, sync_stream=True + streaming_response=None, sync_stream=True, streams_reasoning_signatures=True ).chunk_parser(done) blocks: Final = result.choices[0].delta.thinking_blocks assert blocks is not None @@ -3291,6 +3291,96 @@ def test_reasoning_output_item_done_emits_signed_thinking_blocks_the_stream_keep assert start["data"] == encrypted_reasoning_signature(encrypted) +_SIGNED_REASONING_STREAM_BODY: Final[Mapping[str, object]] = { + "id": "resp_signed_reasoning", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-6-sol", + "output": [ + {"type": "reasoning", "id": "rs_signed", "summary": [], "encrypted_content": "gAAAA_consumer"}, + { + "type": "message", + "id": "msg_logged_call_type", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "pong", "annotations": []}], + }, + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, +} + +_GET_WEATHER_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} + + +def _signed_reasoning_sse() -> str: + events: Final = ( + { + "type": "response.created", + "response": {**_SIGNED_REASONING_STREAM_BODY, "status": "in_progress", "output": [], "usage": None}, + }, + {"type": "response.output_item.done", "output_index": 0, "item": _SIGNED_REASONING_STREAM_BODY["output"][0]}, + { + "type": "response.output_text.delta", + "item_id": "msg_logged_call_type", + "output_index": 1, + "content_index": 0, + "delta": "pong", + }, + {"type": "response.completed", "response": _SIGNED_REASONING_STREAM_BODY}, + ) + return "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events) + + +_AZURE_BRIDGE_KWARGS: Final[Mapping[str, object]] = { + "api_base": "https://example.invalid", + "api_version": "2025-04-01-preview", + "api_key": "test-key", +} + + +@pytest.mark.asyncio +async def test_streamed_reasoning_signature_reaches_messages_clients_and_never_chat_completions_deltas( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock(assert_all_called=True) as router: + router.post(url__startswith="https://example.invalid/openai/responses").mock( + side_effect=lambda _request: httpx.Response( + 200, text=_signed_reasoning_sse(), headers={"content-type": "text/event-stream"} + ) + ) + chat_stream: Final = await litellm.acompletion( + model="azure/gpt-6-sol", + messages=[{"role": "user", "content": "ping"}], + tools=[_GET_WEATHER_TOOL], + stream=True, + **_AZURE_BRIDGE_KWARGS, + ) + chat_deltas: Final = [chunk.choices[0].delta async for chunk in chat_stream if chunk.choices] + messages_stream: Final = await litellm.anthropic.messages.acreate( + model="azure/gpt-6-sol", + messages=[{"role": "user", "content": "ping"}], + tools=[{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}], + max_tokens=64, + stream=True, + **_AZURE_BRIDGE_KWARGS, + ) + sse: Final = b"".join([chunk async for chunk in messages_stream]).decode() + + assert [getattr(delta, "thinking_blocks", None) for delta in chat_deltas if getattr(delta, "thinking_blocks", None)] == [] + assert "".join(delta.content or "" for delta in chat_deltas) == "pong" + messages_events: Final = [json.loads(line[len("data: ") :]) for line in sse.splitlines() if line.startswith("data: ")] + started: Final = [event["content_block"] for event in messages_events if event.get("type") == "content_block_start"] + assert [block["type"] for block in started] == ["redacted_thinking", "text"] + assert started[0]["data"] == encrypted_reasoning_signature("gAAAA_consumer") + + def test_reasoning_output_item_done_with_summary_signs_the_thinking_text(): from litellm.completion_extras.litellm_responses_transformation.transformation import ( OpenAiResponsesToChatCompletionStreamIterator, @@ -3309,7 +3399,7 @@ def test_reasoning_output_item_done_with_summary_signs_the_thinking_text(): }, } result: Final = OpenAiResponsesToChatCompletionStreamIterator( - streaming_response=None, sync_stream=True + streaming_response=None, sync_stream=True, streams_reasoning_signatures=True ).chunk_parser(done) block: Final = result.choices[0].delta.thinking_blocks[0] assert block["type"] == "thinking" @@ -3343,7 +3433,9 @@ def test_streamed_reasoning_summary_reaches_messages_clients_once_with_its_signa }, {"type": "response.output_text.delta", "output_index": 1, "delta": "Oslo."}, ) - parser: Final = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + parser: Final = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True, streams_reasoning_signatures=True + ) chunks: Final = ( *(parser.chunk_parser(event) for event in events), ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")]), @@ -3371,7 +3463,9 @@ def test_reasoning_item_repeated_at_stream_end_reaches_messages_clients_once(): {"type": "response.output_item.done", "output_index": 0, "item": reasoning_item}, {"type": "response.output_text.delta", "output_index": 1, "delta": "Oslo."}, ) - parser: Final = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + parser: Final = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True, streams_reasoning_signatures=True + ) chunks: Final = ( *(parser.chunk_parser(event) for event in events), ModelResponseStream( @@ -3401,7 +3495,7 @@ def test_reasoning_output_item_done_without_encrypted_content_stays_empty(): }, } result: Final = OpenAiResponsesToChatCompletionStreamIterator( - streaming_response=None, sync_stream=True + streaming_response=None, sync_stream=True, streams_reasoning_signatures=True ).chunk_parser(done) assert getattr(result.choices[0].delta, "thinking_blocks", None) in (None, []) diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index cda53b82931..ce7ff6d6eb0 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -5585,6 +5585,62 @@ async def test_messages_do_not_retry_other_rejections( assert route.call_count == 1 +_STRIPPED_CHAT_MODEL: Final = "azure/gpt-5.6" +_STRIPPED_CHAT_API_BASE: Final = "https://stripped-reasoning.openai.azure.com" +_STRIPPED_CHAT_URL: Final = r".*stripped-reasoning\.openai\.azure\.com/openai/deployments/gpt-5\.6/chat/completions.*" + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_messages_do_not_retry_when_the_translated_request_already_dropped_the_reasoning( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=_STRIPPED_CHAT_URL).mock(return_value=_RETRY_REJECTION) + + with pytest.raises(litellm.BadRequestError): + await litellm.anthropic.messages.acreate( + model=_STRIPPED_CHAT_MODEL, + api_base=_STRIPPED_CHAT_API_BASE, + api_key="fake-key", + max_tokens=64, + messages=_history_with_foreign_reasoning(), + ) + + assert route.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_xai_web_search_turn_replays_the_reasoning_its_responses_route_minted( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=r".*xai-search\.example/v1/responses").mock( + return_value=httpx.Response(200, json=_RETRY_REPLY) + ) + + await litellm.anthropic.messages.acreate( + model="xai/grok-4.7", + api_base="https://xai-search.example/v1", + api_key="fake-key", + max_tokens=2048, + tools=[{"type": "web_search_20250305", "name": "web_search"}], + messages=[ + {"role": "user", "content": "Which city is colder today?"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAA_same_org")}, + {"type": "text", "text": "Oslo."}, + ], + }, + {"role": "user", "content": "And by how much?"}, + ], + ) + + assert route.call_count == 1 + assert [item["encrypted_content"] for item in _sent_reasoning_items(route.calls[0].request)] == ["gAAA_same_org"] + + @pytest.mark.asyncio @pytest.mark.usefixtures("_chat_bridge_httpx") async def test_openai_thinking_on_the_chat_adapter_keeps_reasoning_its_responses_route_replays( From d6625245b49a03345a2bea9cd8d5719835a268db Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 01:53:19 -0800 Subject: [PATCH 24/41] fix(anthropic): send openai/ Messages turns again once without encrypted reasoning a backend rejects The openai/ Responses path on /v1/messages had no retry, so a chatgpt/ block left in history answered 400 on every later openai/ turn. That path now sends the turn once more without encrypted reasoning, the same way the chat bridge does. The retry matches only invalid_encrypted_content or a "could not be verified" message naming the content, so the exception mapper's own hint no longer triggers it. The xAI web search keep path is backed out, since xAI sends no encrypted reasoning over the bridge. Backport of 15e60170124 from #49076 to lts/1.108.x, ported by hand onto the pass_through modules. Co-authored-by: DDT <141296842+Hugo-DDT@users.noreply.github.com> --- .../pass_through/adapters/handler.py | 51 ++++-- .../pass_through/adapters/transformation.py | 4 - .../responses_adapters/handler.py | 152 +++++++++++++----- litellm/responses/bridge_check.py | 2 - ...al_pass_through_adapters_transformation.py | 39 +---- .../test_responses_adapters_handler.py | 141 +++++++++++++++- 6 files changed, 294 insertions(+), 95 deletions(-) diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 739077d3566..9cd08aca1d6 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -1,3 +1,4 @@ +import re from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Mapping, Sequence from dataclasses import dataclass from functools import partial @@ -77,6 +78,7 @@ class _CompletionKwargs(TypedDict, total=False, extra_items=object): _ANTHROPIC_MESSAGES: Final = TypeAdapter(list[dict[str, object]]) +_UNVERIFIED_ENCRYPTED_CONTENT: Final = re.compile(r"encrypted content \S+ could not be verified", re.IGNORECASE) _PrepareReplay: TypeAlias = "Callable[[_AnthropicMessages], _CompletionKwargs]" @@ -125,24 +127,49 @@ def _sends_encrypted_reasoning(completion_kwargs: _CompletionKwargs) -> bool: return any(_message_carries_encrypted_reasoning(message) for message in sent) +def _rejects_encrypted_reasoning(error: litellm.BadRequestError) -> bool: + error_message: Final = str(error) + return ( + "invalid_encrypted_content" in error_message or _UNVERIFIED_ENCRYPTED_CONTENT.search(error_message) is not None + ) + + +def messages_without_rejected_reasoning( + messages: Sequence[Mapping[str, object]], + error: litellm.BadRequestError, + *, + sent_encrypted_reasoning: bool, + provider: str, + model: str, +) -> "_AnthropicMessages | None": + if not sent_encrypted_reasoning or not _rejects_encrypted_reasoning(error): + return None + verbose_logger.warning( + "Backend rejected replayed encrypted reasoning, retrying once without it. " + "Enable the encrypted_content_affinity pre-call check to drop such blocks before the call. " + "provider=%s model=%s status=%s code=%s", + provider, + model, + error.status_code, + _logged_rejection_code(error), + ) + return _ANTHROPIC_MESSAGES.validate_python( + strip_encrypted_reasoning_blocks_from_anthropic_messages([dict(message) for message in messages]) + ) + + def _messages_without_rejected_reasoning( messages: _AnthropicMessages, error: litellm.BadRequestError, completion_kwargs: _CompletionKwargs, ) -> "_AnthropicMessages | None": - error_message: Final = str(error) - rejected: Final = "invalid_encrypted_content" in error_message or "could not be verified" in error_message - if not rejected or not _sends_encrypted_reasoning(completion_kwargs): - return None - verbose_logger.warning( - "Backend rejected replayed encrypted reasoning, retrying once without it. " - "provider=%s model=%s status=%s code=%s", - _completion_text(completion_kwargs, "custom_llm_provider"), - _completion_text(completion_kwargs, "model"), - error.status_code, - _logged_rejection_code(error), + return messages_without_rejected_reasoning( + messages, + error, + sent_encrypted_reasoning=_sends_encrypted_reasoning(completion_kwargs), + provider=_completion_text(completion_kwargs, "custom_llm_provider"), + model=_completion_text(completion_kwargs, "model"), ) - return _ANTHROPIC_MESSAGES.validate_python(strip_encrypted_reasoning_blocks_from_anthropic_messages(messages)) async def _acompletion_without_rejected_reasoning( diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index 0f0f7cd22f0..ee553749c74 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -201,7 +201,6 @@ from litellm.types.llms.openai import ( ChatCompletionToolParamFunctionChunk, ChatCompletionToolReferenceObject, ChatCompletionUserMessage, - OpenAIWebSearchOptions, ToolMessageContentPart, ) from litellm.types.utils import Choices, ModelResponse, StreamingChoices, Usage @@ -229,7 +228,6 @@ class ChatRouteInputs: reasoning_effort: str | Mapping[str, object] | None = None api_base: str | None = None thinking: Mapping[str, object] | None = None - web_search: bool = False _NO_ROUTE_INPUTS: Final = ChatRouteInputs() @@ -274,7 +272,6 @@ def _keeps_encrypted_reasoning_for_responses( tools=route.tools, reasoning_effort=route.reasoning_effort, api_base=route.api_base, - web_search_options=OpenAIWebSearchOptions() if route.web_search else None, ) @@ -1352,7 +1349,6 @@ class LiteLLMAnthropicMessagesAdapter: reasoning_effort=new_kwargs.get("reasoning_effort"), api_base=api_base, thinking=anthropic_message_request.get("thinking"), - web_search="web_search_options" in new_kwargs, ), ) ## ADD SYSTEM MESSAGE TO MESSAGES diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py index 3eeaee8ad43..b705216eed4 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py @@ -14,9 +14,10 @@ from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, supports_openai_prompt_cache_breakpoint, ) -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items, normalize_drop_params from litellm.litellm_core_utils.prompt_templates.common_utils import cache_control_as_prompt_cache_breakpoint from litellm.litellm_core_utils.tool_fields import drops_mcp_servers, tools_without_mcp_servers +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS from litellm.types.llms.anthropic import ( AllAnthropicMessageValues, @@ -31,7 +32,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( from litellm.types.llms.openai import ResponsesAPIResponse from litellm.utils import ProviderConfigManager -from ..adapters.handler import request_safeguards_evaluator +from ..adapters.handler import messages_without_rejected_reasoning, request_safeguards_evaluator from ..safeguards import with_safeguard_results from ..utils import litellm_logging_obj_from_kwargs, local_model_name from .streaming_iterator import AnthropicResponsesStreamWrapper @@ -254,6 +255,69 @@ def _build_responses_kwargs( return responses_kwargs +_ResponsesKwargsFor: TypeAlias = Callable[[AnthropicRequestMessages], Mapping[str, object]] +_ResponsesResult: TypeAlias = ResponsesAPIResponse | BaseResponsesAPIStreamingIterator +_SyncResponsesResult: TypeAlias = _ResponsesResult | Coroutine[object, object, _ResponsesResult] + + +def _is_encrypted_reasoning_item(item: object) -> bool: + mapping: Final = as_str_mapping(item) + return mapping is not None and mapping.get("type") == "reasoning" and bool(mapping.get("encrypted_content")) + + +def _sends_encrypted_reasoning(responses_kwargs: Mapping[str, object]) -> bool: + input_items: Final = responses_kwargs.get("input") + sent: Final = iter_items(input_items) if input_items is not None else () + return any(map(_is_encrypted_reasoning_item, sent)) + + +def _kwarg_text(responses_kwargs: Mapping[str, object], key: str) -> str: + value: Final = responses_kwargs.get(key) + return value if isinstance(value, str) else "" + + +def _messages_without_rejected_reasoning( + messages: AnthropicRequestMessages, + error: litellm.BadRequestError, + responses_kwargs: Mapping[str, object], +) -> AnthropicRequestMessages | None: + return messages_without_rejected_reasoning( + messages, + error, + sent_encrypted_reasoning=_sends_encrypted_reasoning(responses_kwargs), + provider=_kwarg_text(responses_kwargs, "custom_llm_provider"), + model=_kwarg_text(responses_kwargs, "model"), + ) + + +async def _aresponses_without_rejected_reasoning( + messages: AnthropicRequestMessages, responses_kwargs_for: _ResponsesKwargsFor +) -> tuple[Mapping[str, object], _ResponsesResult]: + responses_kwargs: Final = responses_kwargs_for(messages) + try: + return responses_kwargs, await litellm.aresponses(**responses_kwargs) + except litellm.BadRequestError as error: + replayed_messages: Final = _messages_without_rejected_reasoning(messages, error, responses_kwargs) + if replayed_messages is None: + raise + replay_kwargs: Final = responses_kwargs_for(replayed_messages) + return replay_kwargs, await litellm.aresponses(**replay_kwargs) + + +def _responses_without_rejected_reasoning( + messages: AnthropicRequestMessages, responses_kwargs_for: _ResponsesKwargsFor +) -> tuple[Mapping[str, object], _SyncResponsesResult]: + responses_kwargs: Final = responses_kwargs_for(messages) + try: + return responses_kwargs, litellm.responses(**responses_kwargs) + except litellm.BadRequestError as error: + replayed_messages: Final = _messages_without_rejected_reasoning(messages, error, responses_kwargs) + if replayed_messages is None: + raise + replay_kwargs: Final = responses_kwargs_for(replayed_messages) + return replay_kwargs, litellm.responses(**replay_kwargs) + + class LiteLLMMessagesToResponsesAPIHandler: """ Handles Anthropic /v1/messages requests for OpenAI / Azure models by @@ -281,28 +345,30 @@ class LiteLLMMessagesToResponsesAPIHandler: output_format: AnthropicOutputSchema | None = None, **kwargs: object, ) -> AnthropicMessagesResponse | AsyncIterator[bytes]: - responses_kwargs: Final = _build_responses_kwargs( - max_tokens=max_tokens, - messages=messages, - model=model, - context_management=context_management, - metadata=metadata, - output_config=output_config, - stop_sequences=stop_sequences, - stream=stream, - system=system, - temperature=temperature, - thinking=thinking, - tool_choice=tool_choice, - tools=tools, - top_k=top_k, - top_p=top_p, - output_format=output_format, - extra_kwargs=kwargs, - ) + + def responses_kwargs_for(sent_messages: AnthropicRequestMessages) -> Mapping[str, object]: + return _build_responses_kwargs( + max_tokens=max_tokens, + messages=sent_messages, + model=model, + context_management=context_management, + metadata=metadata, + output_config=output_config, + stop_sequences=stop_sequences, + stream=stream, + system=system, + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + output_format=output_format, + extra_kwargs=kwargs, + ) safeguards_evaluator: Final = request_safeguards_evaluator(kwargs, messages) - result: Final = await litellm.aresponses(**responses_kwargs) + responses_kwargs, result = await _aresponses_without_rejected_reasoning(messages, responses_kwargs_for) if stream: wrapper: Final = AnthropicResponsesStreamWrapper( @@ -366,28 +432,28 @@ class LiteLLMMessagesToResponsesAPIHandler: **kwargs, ) - # Sync path - responses_kwargs: Final = _build_responses_kwargs( - max_tokens=max_tokens, - messages=messages, - model=model, - context_management=context_management, - metadata=metadata, - output_config=output_config, - stop_sequences=stop_sequences, - stream=stream, - system=system, - temperature=temperature, - thinking=thinking, - tool_choice=tool_choice, - tools=tools, - top_k=top_k, - top_p=top_p, - output_format=output_format, - extra_kwargs=kwargs, - ) + def responses_kwargs_for(sent_messages: AnthropicRequestMessages) -> Mapping[str, object]: + return _build_responses_kwargs( + max_tokens=max_tokens, + messages=sent_messages, + model=model, + context_management=context_management, + metadata=metadata, + output_config=output_config, + stop_sequences=stop_sequences, + stream=stream, + system=system, + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + output_format=output_format, + extra_kwargs=kwargs, + ) - result: Final = litellm.responses(**responses_kwargs) + responses_kwargs, result = _responses_without_rejected_reasoning(messages, responses_kwargs_for) if stream: wrapper: Final = AnthropicResponsesStreamWrapper( diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index 4a25bc815ba..d8a9440bc57 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -198,7 +198,6 @@ def chat_completion_ends_in_responses_api( tools: Sequence[Mapping[str, object]] | None = None, reasoning_effort: str | Mapping[str, object] | None = None, api_base: str | None = None, - web_search_options: OpenAIWebSearchOptions | None = None, ) -> bool: """Whether chat completion for this model is served by a Responses API. @@ -218,7 +217,6 @@ def chat_completion_ends_in_responses_api( model_info, _resolved = responses_api_bridge_check( model=resolved_model, custom_llm_provider=provider, - web_search_options=web_search_options, tools=tools, reasoning_effort=reasoning_effort, api_base=api_base, diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index ce7ff6d6eb0..595f9bb9837 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -5449,6 +5449,9 @@ _RETRY_REJECTION: Final = httpx.Response( _UNRELATED_REJECTION: Final = httpx.Response( 400, json={"error": {"message": "max_output_tokens is too small", "type": "invalid_request_error"}} ) +_UNRELATED_UNVERIFIED_REJECTION: Final = httpx.Response( + 400, json={"error": {"message": "The organization could not be verified.", "type": "invalid_request_error"}} +) _RETRY_REPLY: Final = { "id": "resp_retry", "object": "response", @@ -5523,6 +5526,7 @@ async def test_async_messages_retry_once_without_reasoning_the_backend_cannot_de assert [block["text"] for block in response["content"] if block["type"] == "text"] == ["Oslo"] assert retry_warnings == ( "Backend rejected replayed encrypted reasoning, retrying once without it. " + "Enable the encrypted_content_affinity pre-call check to drop such blocks before the call. " "provider=azure model=azure/gpt-5.4-nano status=400 code=invalid_encrypted_content", ) assert all( @@ -5564,6 +5568,9 @@ def test_sync_messages_retry_once_without_reasoning_the_backend_cannot_decrypt( id="no_encrypted_reasoning_to_drop", ), pytest.param(_history_with_foreign_reasoning(), _UNRELATED_REJECTION, id="unrelated_rejection"), + pytest.param( + _history_with_foreign_reasoning(), _UNRELATED_UNVERIFIED_REJECTION, id="unrelated_unverified_rejection" + ), ], ) async def test_messages_do_not_retry_other_rejections( @@ -5609,38 +5616,6 @@ async def test_messages_do_not_retry_when_the_translated_request_already_dropped assert route.call_count == 1 -@pytest.mark.asyncio -@pytest.mark.usefixtures("_chat_bridge_httpx") -async def test_xai_web_search_turn_replays_the_reasoning_its_responses_route_minted( - respx_mock: respx.MockRouter, -) -> None: - route: Final = respx_mock.post(url__regex=r".*xai-search\.example/v1/responses").mock( - return_value=httpx.Response(200, json=_RETRY_REPLY) - ) - - await litellm.anthropic.messages.acreate( - model="xai/grok-4.7", - api_base="https://xai-search.example/v1", - api_key="fake-key", - max_tokens=2048, - tools=[{"type": "web_search_20250305", "name": "web_search"}], - messages=[ - {"role": "user", "content": "Which city is colder today?"}, - { - "role": "assistant", - "content": [ - {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAA_same_org")}, - {"type": "text", "text": "Oslo."}, - ], - }, - {"role": "user", "content": "And by how much?"}, - ], - ) - - assert route.call_count == 1 - assert [item["encrypted_content"] for item in _sent_reasoning_items(route.calls[0].request)] == ["gAAA_same_org"] - - @pytest.mark.asyncio @pytest.mark.usefixtures("_chat_bridge_httpx") async def test_openai_thinking_on_the_chat_adapter_keeps_reasoning_its_responses_route_replays( diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py index 6aec78bdb97..042f99be001 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py @@ -2,10 +2,11 @@ import datetime import json import os import sys -from collections.abc import Mapping -from typing import Final +from collections.abc import AsyncIterator, Mapping +from typing import Final, cast from unittest.mock import AsyncMock, patch +import httpx import pytest import respx @@ -15,6 +16,7 @@ import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, + encrypted_reasoning_signature, ) from litellm.llms.anthropic.pass_through.responses_adapters.handler import ( LiteLLMMessagesToResponsesAPIHandler, @@ -684,3 +686,138 @@ def test_drop_params_drops_the_mcp_toolset_tools_of_dropped_mcp_servers(monkeypa ) assert [tool["name"] for tool in responses_kwargs["tools"]] == ["get_weather"] + +TEXT_RESPONSE: Final = { + "id": "resp_replayed", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6-luna", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "done", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 3, "output_tokens": 4, "total_tokens": 7}, +} + +REJECTED_REPLAYED_REASONING: Final = { + "error": { + "message": "The encrypted content gAAA_foreign could not be verified.", + "type": "invalid_request_error", + "param": None, + "code": "invalid_encrypted_content", + } +} +UNVERIFIED_ORGANIZATION: Final = { + "error": {"message": "The organization could not be verified.", "type": "invalid_request_error"} +} + + +def _history(*, replayed_reasoning: bool) -> list[dict[str, object]]: + reasoning: Final = ( + ({"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAA_foreign")},) + if replayed_reasoning + else () + ) + return [ + {"role": "user", "content": "List the logs."}, + {"role": "assistant", "content": [*reasoning, {"type": "tool_use", "id": "call_1", "name": "Bash", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "app.log"}]}, + ] + + +def _replayed_reasoning_counts(upstream: respx.Route) -> list[int]: + return [ + sum(item.get("type") == "reasoning" for item in json.loads(call.request.content)["input"]) + for call in upstream.calls + ] + + +def _upstream(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, first: httpx.Response, stream: bool): + monkeypatch.setenv("OPENAI_API_KEY", "sk-replay-test") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + served: Final = ( + httpx.Response(200, headers={"Content-Type": "text/event-stream"}, content=RESPONSES_SSE_BODY) + if stream + else httpx.Response(200, json=TEXT_RESPONSE) + ) + return respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=[first, served]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_openai_turn_is_sent_again_once_without_replayed_reasoning_the_backend_cannot_decrypt( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stream: bool +) -> None: + upstream: Final = _upstream(respx_mock, monkeypatch, httpx.Response(400, json=REJECTED_REPLAYED_REASONING), stream) + + response: Final = await LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + stream=stream, + custom_llm_provider="openai", + ) + events: Final = ( + [json.loads(chunk.decode().split("data: ", 1)[1]) async for chunk in cast(AsyncIterator[bytes], response)] + if stream + else [cast(Mapping[str, object], response)] + ) + + assert _replayed_reasoning_counts(upstream) == [1, 0] + assert events[-1]["type"] in ("message", "message_stop") + + +def test_sync_openai_turn_is_sent_again_once_without_replayed_reasoning_the_backend_cannot_decrypt( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + upstream: Final = _upstream(respx_mock, monkeypatch, httpx.Response(400, json=REJECTED_REPLAYED_REASONING), False) + + response: Final = LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + custom_llm_provider="openai", + ) + + assert _replayed_reasoning_counts(upstream) == [1, 0] + assert cast(Mapping[str, object], response)["content"] == [{"type": "text", "text": "done"}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "replayed_reasoning, rejection", + [ + pytest.param(False, REJECTED_REPLAYED_REASONING, id="no_replayed_reasoning"), + pytest.param(True, UNVERIFIED_ORGANIZATION, id="unrelated_rejection"), + ], +) +async def test_openai_turn_is_not_sent_again_for_other_rejections( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + replayed_reasoning: bool, + rejection: Mapping[str, object], +) -> None: + upstream: Final = _upstream(respx_mock, monkeypatch, httpx.Response(400, json=rejection), False) + + with pytest.raises(litellm.BadRequestError): + await LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( + max_tokens=64, + messages=_history(replayed_reasoning=replayed_reasoning), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + custom_llm_provider="openai", + ) + + assert upstream.call_count == 1 From ce35f27bfaecc0b513b0cd5921b095d80b406a83 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 04:08:08 -0800 Subject: [PATCH 25/41] fix(anthropic): keep every non-streamed reasoning item for Messages clients and read a deployment's reasoning_summary in the keep check Backport of #49076 to lts/1.108.x. Ported from cadfe1d1798 (cursor/chatgpt-encrypted-reasoning-messages-1723) onto the line's pass_through/adapters layout. --- .../handler.py | 18 +++-- .../transformation.py | 40 +++++++---- .../pass_through/adapters/handler.py | 2 + .../pass_through/adapters/transformation.py | 12 +++- litellm/responses/bridge_check.py | 2 + ...responses_transformation_transformation.py | 66 +++++++++++++++++++ ...al_pass_through_adapters_transformation.py | 28 ++++++++ 7 files changed, 144 insertions(+), 24 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index eef9cc6c4c9..e824d9bcfc1 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -13,9 +13,7 @@ from typing import ( from typing_extensions import TypedDict -from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.types.llms.openai import ResponsesAPIResponse -from litellm.types.utils import CallTypes if TYPE_CHECKING: from litellm import CustomStreamWrapper, LiteLLMLoggingObj, ModelResponse @@ -38,10 +36,6 @@ def _restore_routing_prefix(model: str, custom_llm_provider: str) -> str: return f"{custom_llm_provider}/{model}" -def _serves_anthropic_messages(logging_obj: "LiteLLMLoggingObj") -> bool: - return attribute_or_none(logging_obj, "call_type") == CallTypes.anthropic_messages.value - - class ResponsesToCompletionBridgeHandler: def __init__(self): from .transformation import LiteLLMResponsesTransformationHandler @@ -188,7 +182,9 @@ class ResponsesToCompletionBridgeHandler: model_response: Final = validated_kwargs["model_response"] logging_obj: Final = validated_kwargs["logging_obj"] custom_llm_provider: Final = validated_kwargs["custom_llm_provider"] - serves_anthropic_messages: Final = _serves_anthropic_messages(logging_obj) + from .transformation import serves_anthropic_messages + + streams_reasoning_signatures: Final = serves_anthropic_messages(logging_obj) if kwargs.get("stream") is True and "stream" not in optional_params: optional_params = {**optional_params, "stream": True} @@ -267,7 +263,7 @@ class ResponsesToCompletionBridgeHandler: structured_output=is_structured_output_request( cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict ), - streams_reasoning_signatures=serves_anthropic_messages, + streams_reasoning_signatures=streams_reasoning_signatures, ) streamwrapper: Final = CustomStreamWrapper( completion_stream=completion_stream, @@ -290,7 +286,9 @@ class ResponsesToCompletionBridgeHandler: model_response: Final = validated_kwargs["model_response"] logging_obj: Final = validated_kwargs["logging_obj"] custom_llm_provider: Final = validated_kwargs["custom_llm_provider"] - serves_anthropic_messages: Final = _serves_anthropic_messages(logging_obj) + from .transformation import serves_anthropic_messages + + streams_reasoning_signatures: Final = serves_anthropic_messages(logging_obj) if kwargs.get("stream") is True and "stream" not in optional_params: optional_params = {**optional_params, "stream": True} @@ -372,7 +370,7 @@ class ResponsesToCompletionBridgeHandler: structured_output=is_structured_output_request( cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict ), - streams_reasoning_signatures=serves_anthropic_messages, + streams_reasoning_signatures=streams_reasoning_signatures, ) streamwrapper: Final = CustomStreamWrapper( completion_stream=completion_stream, diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 33ac222b06b..cf8d845bf26 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -26,7 +26,7 @@ import litellm from litellm import ModelResponse from litellm._logging import verbose_logger from litellm.integrations.anthropic_cache_control_hook import supports_openai_prompt_cache_breakpoint -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( responses_reasoning_items_from_thinking_blocks, @@ -58,7 +58,7 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ResponsesAPIStreamEvents, ) -from litellm.types.utils import Delta, GenericStreamingChunk, ModelResponseStream, StreamingChoices +from litellm.types.utils import CallTypes, Delta, GenericStreamingChunk, ModelResponseStream, StreamingChoices if TYPE_CHECKING: from openai.types.responses import ResponseInputImageParam, ResponseOutputItem @@ -227,6 +227,10 @@ class _BuiltReasoningItem(TypedDict): summary: Sequence[_ReasoningSummaryText] +def serves_anthropic_messages(logging_obj: object) -> bool: + return attribute_or_none(logging_obj, "call_type") == CallTypes.anthropic_messages.value + + def _get_reasoning_items( msg: "AllMessageValues", ) -> list[ChatCompletionReasoningItem]: @@ -549,7 +553,7 @@ def _fold_choices_into_one( choices: Sequence["Choices"], tool_calls: Sequence[Mapping[str, object]], trailing_reasoning_content: str | None, - trailing_reasoning_item: _BuiltReasoningItem | None, + trailing_reasoning_items: Sequence[_BuiltReasoningItem], tool_search_items: Sequence[Mapping[str, object]], structured_output: bool = False, ) -> "Choices": @@ -568,7 +572,7 @@ def _fold_choices_into_one( reasoning_contents: Final = (*(_message_reasoning_content(m) for m in messages), trailing_reasoning_content) reasoning_items: Final = ( *chain.from_iterable(_message_reasoning_items(m) for m in messages), - *(() if trailing_reasoning_item is None else (trailing_reasoning_item,)), + *trailing_reasoning_items, ) return Choices( message=Message( @@ -1222,6 +1226,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): output_items: Sequence[object], handle_raw_dict_callback: Callable[..., tuple["Choices | None", int]] | None = None, structured_output: bool = False, + keeps_every_reasoning_item: bool = False, ) -> list["Choices"]: """ Convert Responses API output items to chat completion choices. @@ -1231,6 +1236,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): handle_raw_dict_callback: Optional callback for handling raw dict items structured_output: The chat request carried a json response_format, so only the first message item becomes content + keeps_every_reasoning_item: Keep each reasoning item before an answer, not only the last Returns: List of Choices objects @@ -1253,7 +1259,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): choices: Final[list[Choices]] = [] index = 0 reasoning_content: str | None = None - pending_reasoning_item: _BuiltReasoningItem | None = None + pending_reasoning_items: tuple[_BuiltReasoningItem, ...] = () # Collect all tool calls to put them in a single choice # (Chat Completions API expects all tool calls in one message) @@ -1262,12 +1268,19 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): for item in output_items: if isinstance(item, ResponseReasoningItem): - pending_reasoning_item = _build_reasoning_item( - item_id=item.id, - encrypted_content=getattr(item, "encrypted_content", None), - summary_raw=item.summary, + pending_reasoning_items = ( + *(pending_reasoning_items if keeps_every_reasoning_item else ()), + _build_reasoning_item( + item_id=item.id, + encrypted_content=getattr(item, "encrypted_content", None), + summary_raw=item.summary, + ), + ) + reasoning_content = " ".join( + s["text"] + for s in chain.from_iterable(built["summary"] for built in pending_reasoning_items) + if s.get("text") ) - reasoning_content = " ".join(s["text"] for s in pending_reasoning_item["summary"] if s.get("text")) elif isinstance(item, ResponseOutputMessage): for content in item.content: @@ -1284,7 +1297,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): annotations=annotations, reasoning_items=cast( list[ChatCompletionReasoningItem] | None, - ([pending_reasoning_item] if pending_reasoning_item is not None else None), + (list(pending_reasoning_items) or None), ), ) @@ -1297,7 +1310,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) reasoning_content = None # flush - pending_reasoning_item = None # flush + pending_reasoning_items = () # flush index += 1 elif isinstance(item, ResponseFunctionToolCall): @@ -1362,7 +1375,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): choices=choices, tool_calls=accumulated_tool_calls, trailing_reasoning_content=reasoning_content, - trailing_reasoning_item=pending_reasoning_item, + trailing_reasoning_items=pending_reasoning_items, tool_search_items=tool_search_items, structured_output=structured_output, ) @@ -1500,6 +1513,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): structured_output=is_structured_output_request( cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict ), + keeps_every_reasoning_item=serves_anthropic_messages(logging_obj), ) response_is_incomplete: Final = raw_response.status == "incomplete" or ( diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 9cd08aca1d6..4f21c7dbe5e 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -17,6 +17,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases from litellm.litellm_core_utils.prompt_templates.common_utils import ( is_encrypted_reasoning_block, thinking_blocks_with_encrypted_reasoning, @@ -760,6 +761,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: request_data, custom_llm_provider=target_provider, api_base=api_base if isinstance(api_base, str) else None, + reasoning_summary=peek_reasoning_summary_aliases(provider_kwargs), ) if openai_request is None: diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index ee553749c74..17e50498d46 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -228,6 +228,7 @@ class ChatRouteInputs: reasoning_effort: str | Mapping[str, object] | None = None api_base: str | None = None thinking: Mapping[str, object] | None = None + reasoning_summary: object | None = None _NO_ROUTE_INPUTS: Final = ChatRouteInputs() @@ -271,6 +272,7 @@ def _keeps_encrypted_reasoning_for_responses( custom_llm_provider, tools=route.tools, reasoning_effort=route.reasoning_effort, + reasoning_summary=route.reasoning_summary, api_base=route.api_base, ) @@ -313,7 +315,12 @@ class AnthropicAdapter: return result def translate_completion_input_params_with_tool_mapping( - self, kwargs, *, custom_llm_provider: str | None = None, api_base: str | None = None + self, + kwargs, + *, + custom_llm_provider: str | None = None, + api_base: str | None = None, + reasoning_summary: object | None = None, ) -> tuple[ChatCompletionRequest | None, dict[str, str]]: """ Translate Anthropic request params to OpenAI format, returning tool name mapping. @@ -348,6 +355,7 @@ class AnthropicAdapter: anthropic_message_request=request_body, custom_llm_provider=custom_llm_provider, api_base=api_base, + reasoning_summary=reasoning_summary, ) return translated_body, tool_name_mapping @@ -1295,6 +1303,7 @@ class LiteLLMAnthropicMessagesAdapter: *, custom_llm_provider: str | None = None, api_base: str | None = None, + reasoning_summary: object | None = None, preserve_midturn_system: bool = False, ) -> tuple[ChatCompletionRequest, dict[str, str]]: """ @@ -1349,6 +1358,7 @@ class LiteLLMAnthropicMessagesAdapter: reasoning_effort=new_kwargs.get("reasoning_effort"), api_base=api_base, thinking=anthropic_message_request.get("thinking"), + reasoning_summary=reasoning_summary, ), ) ## ADD SYSTEM MESSAGE TO MESSAGES diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index d8a9440bc57..6b80132a077 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -197,6 +197,7 @@ def chat_completion_ends_in_responses_api( *, tools: Sequence[Mapping[str, object]] | None = None, reasoning_effort: str | Mapping[str, object] | None = None, + reasoning_summary: object | None = None, api_base: str | None = None, ) -> bool: """Whether chat completion for this model is served by a Responses API. @@ -219,6 +220,7 @@ def chat_completion_ends_in_responses_api( custom_llm_provider=provider, tools=tools, reasoning_effort=reasoning_effort, + reasoning_summary=reasoning_summary, api_base=api_base, ) return model_info.get("mode") == "responses" diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index df9d048fb45..cb4cbfc1c57 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3570,6 +3570,72 @@ def test_transform_response_hands_encrypted_reasoning_to_messages_clients_only() assert content[0]["data"] == encrypted_reasoning_signature(encrypted) +@pytest.mark.parametrize( + ("call_type", "kept_ids"), + [ + pytest.param("anthropic_messages", ["rs_first", "rs_second"], id="messages_client_keeps_every_item"), + pytest.param("acompletion", ["rs_second"], id="chat_client_keeps_the_last_item"), + ], +) +def test_non_streamed_tool_turn_hands_messages_clients_every_reasoning_item( + call_type: str, kept_ids: list[str] +) -> None: + from types import SimpleNamespace + from unittest.mock import Mock + + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.utils import ModelResponse, Usage + + def reasoning(item_id: str) -> ResponseReasoningItem: + return ResponseReasoningItem( + id=item_id, summary=[], type="reasoning", content=None, encrypted_content=f"gAAAA_{item_id}", status=None + ) + + def call(index: int, city: str) -> ResponseFunctionToolCall: + return ResponseFunctionToolCall( + id=f"fc_{index}", + type="function_call", + status="completed", + arguments=json.dumps({"city": city}), + call_id=f"call_{index}", + name="get_weather", + ) + + raw: Final = ResponsesAPIResponse.model_construct( + id="resp_two_reasoning", + created_at=1, + error=None, + incomplete_details=None, + model="gpt-5.4-nano", + output=[reasoning("rs_first"), call(1, "Oslo"), reasoning("rs_second"), call(2, "Lyon")], + status="completed", + usage=None, + ) + result: Final = LiteLLMResponsesTransformationHandler().transform_response( + model="gpt-5.4-nano", + raw_response=raw, + model_response=ModelResponse( + id="chatcmpl-two-reasoning", + created=1, + model=None, + object="chat.completion", + system_fingerprint=None, + choices=[], + usage=Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0), + ), + logging_obj=SimpleNamespace(call_type=call_type), + request_data={"model": "gpt-5.4-nano"}, + messages=[{"role": "user", "content": "Oslo or Lyon?"}], + optional_params={}, + litellm_params={"custom_llm_provider": "azure"}, + encoding=Mock(), + ) + + message: Final = result.choices[0].message + assert [item["id"] for item in message.reasoning_items] == kept_ids + assert [tool_call.function.name for tool_call in message.tool_calls] == ["get_weather", "get_weather"] + + def test_streaming_function_call_tool_id_for_degenerate_call_id(): """In streaming, Bedrock Mantle's function_call event carries a unique ``id`` (``fc_...``) and a non-unique, index-based ``call_id`` (``call_0``). For that diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 595f9bb9837..1223185a05c 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -698,6 +698,34 @@ def test_encrypted_reasoning_is_kept_when_the_request_tools_route_it_to_response assert bool(assistant.get("thinking_blocks")) is kept +@pytest.mark.parametrize( + ("deployment_params", "kept"), + [ + pytest.param({"reasoning_summary": "auto"}, True, id="deployment_summary_bridges_to_responses"), + pytest.param({"extra_body": {"reasoningSummary": "auto"}}, True, id="extra_body_summary_bridges_to_responses"), + pytest.param({}, False, id="no_summary_stays_on_chat"), + ], +) +def test_encrypted_reasoning_is_kept_when_the_deployment_reasoning_summary_routes_it_to_responses( + deployment_params: dict[str, object], kept: bool +) -> None: + from litellm.llms.anthropic.pass_through.adapters.handler import ( + LiteLLMMessagesToCompletionTransformationHandler, + ) + + request: Final = _replayed_reasoning_request([]) + completion_kwargs, _mapping = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( + max_tokens=64, + messages=cast(list[dict[str, object]], request["messages"]), + model="azure/gpt-5.1", + thinking={"type": "enabled", "budget_tokens": 2048}, + extra_kwargs={"api_base": "https://keep-check.openai.azure.com", **deployment_params}, + ) + + assistant: Final = cast(list[Mapping[str, object]], completion_kwargs["messages"])[1] + assert bool(assistant.get("thinking_blocks")) is kept + + def test_translate_anthropic_messages_to_openai_sets_reasoning_content(): """Reasoning-aware chat providers read reasoning_content, so thinking text must land there. From 97df7cb369397dc9326fcc59d3be492b79707fe4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 06:06:59 -0800 Subject: [PATCH 26/41] fix(anthropic): keep every non-streamed reasoning item for Messages clients on the line's bridge This line's transform_request rewrites the logging call_type to responses, so the non-streamed reply read the Messages-client check after the rewrite and kept only the last reasoning item. The bridge handler now reads it before the request transform and passes it to transform_response, as the streamed path already does. --- .../handler.py | 12 ++- .../transformation.py | 3 +- ...itellm_responses_transformation_handler.py | 98 +++++++++++++++++++ ...responses_transformation_transformation.py | 12 +-- 4 files changed, 114 insertions(+), 11 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index e824d9bcfc1..c47bb046021 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -184,7 +184,7 @@ class ResponsesToCompletionBridgeHandler: custom_llm_provider: Final = validated_kwargs["custom_llm_provider"] from .transformation import serves_anthropic_messages - streams_reasoning_signatures: Final = serves_anthropic_messages(logging_obj) + serves_messages_client: Final = serves_anthropic_messages(logging_obj) if kwargs.get("stream") is True and "stream" not in optional_params: optional_params = {**optional_params, "stream": True} @@ -225,6 +225,7 @@ class ResponsesToCompletionBridgeHandler: encoding=validated_kwargs["encoding"], api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), + keeps_every_reasoning_item=serves_messages_client, ) elif isinstance(result, ModelResponse): if not stream: @@ -250,6 +251,7 @@ class ResponsesToCompletionBridgeHandler: encoding=validated_kwargs["encoding"], api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), + keeps_every_reasoning_item=serves_messages_client, ) else: if self._is_preformatted_cached_chat_stream(result): @@ -263,7 +265,7 @@ class ResponsesToCompletionBridgeHandler: structured_output=is_structured_output_request( cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict ), - streams_reasoning_signatures=streams_reasoning_signatures, + streams_reasoning_signatures=serves_messages_client, ) streamwrapper: Final = CustomStreamWrapper( completion_stream=completion_stream, @@ -288,7 +290,7 @@ class ResponsesToCompletionBridgeHandler: custom_llm_provider: Final = validated_kwargs["custom_llm_provider"] from .transformation import serves_anthropic_messages - streams_reasoning_signatures: Final = serves_anthropic_messages(logging_obj) + serves_messages_client: Final = serves_anthropic_messages(logging_obj) if kwargs.get("stream") is True and "stream" not in optional_params: optional_params = {**optional_params, "stream": True} @@ -332,6 +334,7 @@ class ResponsesToCompletionBridgeHandler: encoding=validated_kwargs["encoding"], api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), + keeps_every_reasoning_item=serves_messages_client, ) elif isinstance(result, ModelResponse): if not stream: @@ -357,6 +360,7 @@ class ResponsesToCompletionBridgeHandler: encoding=validated_kwargs["encoding"], api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), + keeps_every_reasoning_item=serves_messages_client, ) else: if self._is_preformatted_cached_chat_stream(result): @@ -370,7 +374,7 @@ class ResponsesToCompletionBridgeHandler: structured_output=is_structured_output_request( cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict ), - streams_reasoning_signatures=streams_reasoning_signatures, + streams_reasoning_signatures=serves_messages_client, ) streamwrapper: Final = CustomStreamWrapper( completion_stream=completion_stream, diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index cf8d845bf26..4c15f2cee4a 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -1482,6 +1482,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): encoding: object, api_key: str | None = None, json_mode: bool | None = None, + keeps_every_reasoning_item: bool = False, ) -> "ModelResponse": """Transform Responses API response to chat completion response""" from litellm.responses.utils import ResponseAPILoggingUtils @@ -1513,7 +1514,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): structured_output=is_structured_output_request( cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict ), - keeps_every_reasoning_item=serves_anthropic_messages(logging_obj), + keeps_every_reasoning_item=keeps_every_reasoning_item, ) response_is_incomplete: Final = raw_response.status == "incomplete" or ( diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py index c5d7ca96a21..0899c86cfb7 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py @@ -335,3 +335,101 @@ async def test_acompletion_keeps_provider_native_model_id_through_responses( handed_model = responses_call.call_args.kwargs["model"] assert _upstream_model_for(handed_model, custom_llm_provider) == expected_upstream_model + + +def _two_reasoning_tool_turn(): + from openai.types.responses import ResponseFunctionToolCall + from openai.types.responses.response_reasoning_item import ResponseReasoningItem + + from litellm.types.llms.openai import ResponsesAPIResponse + + def reasoning(item_id: str) -> ResponseReasoningItem: + return ResponseReasoningItem( + id=item_id, summary=[], type="reasoning", encrypted_content=f"gAAAA_{item_id}" + ) + + def call(index: int, city: str) -> ResponseFunctionToolCall: + return ResponseFunctionToolCall( + id=f"fc_{index}", + type="function_call", + status="completed", + arguments=f'{{"city": "{city}"}}', + call_id=f"call_{index}", + name="get_weather", + ) + + return ResponsesAPIResponse.model_construct( + id="resp_two_reasoning", + created_at=1, + error=None, + incomplete_details=None, + model="gpt-5.4-nano", + output=[reasoning("rs_first"), call(1, "Oslo"), reasoning("rs_second"), call(2, "Lyon")], + status="completed", + usage=None, + ) + + +class _CompletedStream: + def __init__(self, response): + self.completed_response = MagicMock(response=response) + self._hidden_params = {} + + def __iter__(self): + return iter(()) + + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + +def _messages_client_kwargs(call_type: str): + messages = [{"role": "user", "content": "Weather in Oslo and Lyon?"}] + return { + "model": "gpt-5.4-nano", + "custom_llm_provider": "azure", + "messages": messages, + "optional_params": {}, + "litellm_params": {}, + "headers": {}, + "model_response": ModelResponse(), + "logging_obj": LiteLLMLogging( + litellm_call_id="test-call", + call_type=call_type, + model="gpt-5.4-nano", + messages=messages, + function_id="fn-id", + stream=False, + start_time=datetime.now(), + ), + } + + +@pytest.mark.parametrize("is_async", [True, False], ids=["acompletion", "completion"]) +@pytest.mark.parametrize("arrives_streamed", [False, True], ids=["response", "collected_stream"]) +@pytest.mark.parametrize( + ("call_type", "kept_ids"), + [ + pytest.param("anthropic_messages", ["rs_first", "rs_second"], id="messages_client"), + pytest.param("acompletion", ["rs_second"], id="chat_client"), + ], +) +@pytest.mark.asyncio +async def test_non_streamed_bridge_hands_messages_clients_every_reasoning_item( + is_async: bool, arrives_streamed: bool, call_type: str, kept_ids: list[str] +): + turn = _two_reasoning_tool_turn() + upstream = _CompletedStream(turn) if arrives_streamed else turn + bridge = ResponsesToCompletionBridgeHandler() + kwargs = _messages_client_kwargs(call_type) + + if is_async: + with patch("litellm.aresponses", new=AsyncMock(return_value=upstream)): + result = await bridge.acompletion(**kwargs) + else: + with patch("litellm.responses", return_value=upstream): + result = bridge.completion(**kwargs) + + assert [item["id"] for item in result.choices[0].message.reasoning_items] == kept_ids diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index cb4cbfc1c57..04505b717af 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3571,16 +3571,15 @@ def test_transform_response_hands_encrypted_reasoning_to_messages_clients_only() @pytest.mark.parametrize( - ("call_type", "kept_ids"), + ("keeps_every_reasoning_item", "kept_ids"), [ - pytest.param("anthropic_messages", ["rs_first", "rs_second"], id="messages_client_keeps_every_item"), - pytest.param("acompletion", ["rs_second"], id="chat_client_keeps_the_last_item"), + pytest.param(True, ["rs_first", "rs_second"], id="messages_client_keeps_every_item"), + pytest.param(False, ["rs_second"], id="chat_client_keeps_the_last_item"), ], ) def test_non_streamed_tool_turn_hands_messages_clients_every_reasoning_item( - call_type: str, kept_ids: list[str] + keeps_every_reasoning_item: bool, kept_ids: list[str] ) -> None: - from types import SimpleNamespace from unittest.mock import Mock from litellm.types.llms.openai import ResponsesAPIResponse @@ -3623,12 +3622,13 @@ def test_non_streamed_tool_turn_hands_messages_clients_every_reasoning_item( choices=[], usage=Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0), ), - logging_obj=SimpleNamespace(call_type=call_type), + logging_obj=Mock(), request_data={"model": "gpt-5.4-nano"}, messages=[{"role": "user", "content": "Oslo or Lyon?"}], optional_params={}, litellm_params={"custom_llm_provider": "azure"}, encoding=Mock(), + keeps_every_reasoning_item=keeps_every_reasoning_item, ) message: Final = result.choices[0].message From a1c02f7d1a24ba1e03d4ec9831dd471ef3ad18a9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 07:33:23 -0800 Subject: [PATCH 27/41] fix(anthropic): log only the outcome of a turn sent again, and keep every reasoning item in Messages logs on the line A /v1/messages turn whose replayed encrypted block the backend rejects is sent once more without it. The rejected first attempt no longer fires failure callbacks or releases its parallel slot and budget reservation, so only the turn sent again is logged and counted. Success logs for a Messages client now keep every reasoning item of a non-streamed turn, matching the reply the client got. --- litellm/_internal_context.py | 18 ++ litellm/litellm_core_utils/litellm_logging.py | 2 + .../pass_through/adapters/handler.py | 29 ++- .../responses_adapters/handler.py | 25 ++- litellm/utils.py | 6 +- ...itellm_responses_transformation_handler.py | 183 ++++++++++-------- ...al_pass_through_adapters_transformation.py | 84 +++++++- .../test_responses_adapters_handler.py | 139 +++++++++++++ 8 files changed, 384 insertions(+), 102 deletions(-) diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index df549701cf4..1b07c017a3c 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -36,6 +36,8 @@ REDIS_FAMILIES_METADATA_KEY: Final = "families" _service_caller: Final[ContextVar[str | None]] = ContextVar("service_caller", default=None) +_resent_on: Final[ContextVar[Callable[[Exception], bool] | None]] = ContextVar("resent_on", default=None) + @contextmanager def post_response_phase() -> Generator[None]: @@ -113,6 +115,22 @@ def current_service_caller() -> str | None: return _service_caller.get() +@contextmanager +def resent_on(rejection: Callable[[Exception], bool]) -> Generator[None]: + """An attempt its caller sends again when it fails with ``rejection``. That failure is not the request's + outcome, so nothing logs it and its budget reservation and rate-limit slot stay with the attempt sent again.""" + token: Final = _resent_on.set(rejection) + try: + yield + finally: + _restore(_resent_on, token) + + +def is_resent_failure(error: Exception) -> bool: + rejection: Final = _resent_on.get() + return rejection is not None and rejection(error) + + @contextmanager def pinned_billing_time(moment: datetime) -> Generator[None]: """Price every rate lookup inside this block at ``moment`` rather than at each one's own clock read.""" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 95323c2ef5e..8cd8e55c292 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4593,6 +4593,7 @@ class Logging(LiteLLMLoggingBaseClass): from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, responses_provider_metadata, + serves_anthropic_messages, ) served_id: Final = _provider_response_id(result) @@ -4607,6 +4608,7 @@ class Logging(LiteLLMLoggingBaseClass): optional_params={}, litellm_params={}, encoding=litellm.encoding, + keeps_every_reasoning_item=serves_anthropic_messages(self), ) translated.id = served_id or translated.id return translated diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 4f21c7dbe5e..a818d6f3565 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -1,5 +1,6 @@ import re from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Mapping, Sequence +from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass from functools import partial from typing import ( @@ -13,6 +14,7 @@ from pydantic import TypeAdapter from typing_extensions import TypedDict import litellm +from litellm._internal_context import resent_on from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items @@ -135,6 +137,14 @@ def _rejects_encrypted_reasoning(error: litellm.BadRequestError) -> bool: ) +def _is_rejected_reasoning(error: Exception) -> bool: + return isinstance(error, litellm.BadRequestError) and _rejects_encrypted_reasoning(error) + + +def resent_without_rejected_reasoning(sent_encrypted_reasoning: bool) -> AbstractContextManager[None]: + return resent_on(_is_rejected_reasoning) if sent_encrypted_reasoning else nullcontext() + + def messages_without_rejected_reasoning( messages: Sequence[Mapping[str, object]], error: litellm.BadRequestError, @@ -163,11 +173,12 @@ def _messages_without_rejected_reasoning( messages: _AnthropicMessages, error: litellm.BadRequestError, completion_kwargs: _CompletionKwargs, + sent_encrypted_reasoning: bool, ) -> "_AnthropicMessages | None": return messages_without_rejected_reasoning( messages, error, - sent_encrypted_reasoning=_sends_encrypted_reasoning(completion_kwargs), + sent_encrypted_reasoning=sent_encrypted_reasoning, provider=_completion_text(completion_kwargs, "custom_llm_provider"), model=_completion_text(completion_kwargs, "model"), ) @@ -178,10 +189,14 @@ async def _acompletion_without_rejected_reasoning( messages: _AnthropicMessages, prepare_replay: _PrepareReplay, ) -> "ChatResult": + sent_encrypted_reasoning: Final = _sends_encrypted_reasoning(completion_kwargs) try: - return await litellm.acompletion(**completion_kwargs) + with resent_without_rejected_reasoning(sent_encrypted_reasoning): + return await litellm.acompletion(**completion_kwargs) except litellm.BadRequestError as error: - replayed_messages: Final = _messages_without_rejected_reasoning(messages, error, completion_kwargs) + replayed_messages: Final = _messages_without_rejected_reasoning( + messages, error, completion_kwargs, sent_encrypted_reasoning + ) if replayed_messages is None: raise return await litellm.acompletion(**prepare_replay(replayed_messages)) @@ -192,10 +207,14 @@ def _completion_without_rejected_reasoning( messages: _AnthropicMessages, prepare_replay: _PrepareReplay, ) -> "ChatResult | Coroutine[object, object, ChatResult]": + sent_encrypted_reasoning: Final = _sends_encrypted_reasoning(completion_kwargs) try: - return litellm.completion(**completion_kwargs) + with resent_without_rejected_reasoning(sent_encrypted_reasoning): + return litellm.completion(**completion_kwargs) except litellm.BadRequestError as error: - replayed_messages: Final = _messages_without_rejected_reasoning(messages, error, completion_kwargs) + replayed_messages: Final = _messages_without_rejected_reasoning( + messages, error, completion_kwargs, sent_encrypted_reasoning + ) if replayed_messages is None: raise return litellm.completion(**prepare_replay(replayed_messages)) diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py index b705216eed4..abe0f2f419d 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py @@ -32,7 +32,11 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( from litellm.types.llms.openai import ResponsesAPIResponse from litellm.utils import ProviderConfigManager -from ..adapters.handler import messages_without_rejected_reasoning, request_safeguards_evaluator +from ..adapters.handler import ( + messages_without_rejected_reasoning, + request_safeguards_evaluator, + resent_without_rejected_reasoning, +) from ..safeguards import with_safeguard_results from ..utils import litellm_logging_obj_from_kwargs, local_model_name from .streaming_iterator import AnthropicResponsesStreamWrapper @@ -280,11 +284,12 @@ def _messages_without_rejected_reasoning( messages: AnthropicRequestMessages, error: litellm.BadRequestError, responses_kwargs: Mapping[str, object], + sent_encrypted_reasoning: bool, ) -> AnthropicRequestMessages | None: return messages_without_rejected_reasoning( messages, error, - sent_encrypted_reasoning=_sends_encrypted_reasoning(responses_kwargs), + sent_encrypted_reasoning=sent_encrypted_reasoning, provider=_kwarg_text(responses_kwargs, "custom_llm_provider"), model=_kwarg_text(responses_kwargs, "model"), ) @@ -294,10 +299,14 @@ async def _aresponses_without_rejected_reasoning( messages: AnthropicRequestMessages, responses_kwargs_for: _ResponsesKwargsFor ) -> tuple[Mapping[str, object], _ResponsesResult]: responses_kwargs: Final = responses_kwargs_for(messages) + sent_encrypted_reasoning: Final = _sends_encrypted_reasoning(responses_kwargs) try: - return responses_kwargs, await litellm.aresponses(**responses_kwargs) + with resent_without_rejected_reasoning(sent_encrypted_reasoning): + return responses_kwargs, await litellm.aresponses(**responses_kwargs) except litellm.BadRequestError as error: - replayed_messages: Final = _messages_without_rejected_reasoning(messages, error, responses_kwargs) + replayed_messages: Final = _messages_without_rejected_reasoning( + messages, error, responses_kwargs, sent_encrypted_reasoning + ) if replayed_messages is None: raise replay_kwargs: Final = responses_kwargs_for(replayed_messages) @@ -308,10 +317,14 @@ def _responses_without_rejected_reasoning( messages: AnthropicRequestMessages, responses_kwargs_for: _ResponsesKwargsFor ) -> tuple[Mapping[str, object], _SyncResponsesResult]: responses_kwargs: Final = responses_kwargs_for(messages) + sent_encrypted_reasoning: Final = _sends_encrypted_reasoning(responses_kwargs) try: - return responses_kwargs, litellm.responses(**responses_kwargs) + with resent_without_rejected_reasoning(sent_encrypted_reasoning): + return responses_kwargs, litellm.responses(**responses_kwargs) except litellm.BadRequestError as error: - replayed_messages: Final = _messages_without_rejected_reasoning(messages, error, responses_kwargs) + replayed_messages: Final = _messages_without_rejected_reasoning( + messages, error, responses_kwargs, sent_encrypted_reasoning + ) if replayed_messages is None: raise replay_kwargs: Final = responses_kwargs_for(replayed_messages) diff --git a/litellm/utils.py b/litellm/utils.py index da619a5aefe..3213e343c4f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -51,7 +51,7 @@ import litellm.litellm_core_utils # audio_utils.utils is lazy-loaded - only imported when needed for transcription calls import litellm.litellm_core_utils.json_validation_rule -from litellm._internal_context import is_internal_call +from litellm._internal_context import is_internal_call, is_resent_failure from litellm._lazy_imports import ( _get_modified_max_tokens, get_default_encoding, @@ -2032,7 +2032,7 @@ def client(original_function): end_time = datetime.datetime.now() # LOG FAILURE - handle streaming failure logging in the _next_ object, remove `handle_failure` once it's deprecated - if logging_obj: + if logging_obj and not is_resent_failure(e): logging_obj.failure_handler( e, traceback_exception, start_time, end_time ) # DO NOT MAKE THREADED - router retry fallback relies on this! @@ -2279,7 +2279,7 @@ def client(original_function): # Reuse the timestamp taken right when the deployment call itself failed, before # the failure hook ran, so a slow callback doesn't inflate the reported duration. end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now() # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with - if logging_obj and not _is_litellm_internal_call: + if logging_obj and not _is_litellm_internal_call and not is_resent_failure(e): unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) try: logging_obj.failure_handler( diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py index 0899c86cfb7..af0a0dbe9b0 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py @@ -1,7 +1,13 @@ +import json +import time from datetime import datetime +from pathlib import Path +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx import litellm @@ -11,6 +17,7 @@ from litellm.completion_extras.litellm_responses_transformation.handler import ( from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.router import GenericLiteLLMParams +from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature from litellm.types.utils import ModelResponse @@ -337,99 +344,105 @@ async def test_acompletion_keeps_provider_native_model_id_through_responses( assert _upstream_model_for(handed_model, custom_llm_provider) == expected_upstream_model -def _two_reasoning_tool_turn(): - from openai.types.responses import ResponseFunctionToolCall - from openai.types.responses.response_reasoning_item import ResponseReasoningItem - - from litellm.types.llms.openai import ResponsesAPIResponse - - def reasoning(item_id: str) -> ResponseReasoningItem: - return ResponseReasoningItem( - id=item_id, summary=[], type="reasoning", encrypted_content=f"gAAAA_{item_id}" - ) - - def call(index: int, city: str) -> ResponseFunctionToolCall: - return ResponseFunctionToolCall( - id=f"fc_{index}", - type="function_call", - status="completed", - arguments=f'{{"city": "{city}"}}', - call_id=f"call_{index}", - name="get_weather", - ) - - return ResponsesAPIResponse.model_construct( - id="resp_two_reasoning", - created_at=1, - error=None, - incomplete_details=None, - model="gpt-5.4-nano", - output=[reasoning("rs_first"), call(1, "Oslo"), reasoning("rs_second"), call(2, "Lyon")], - status="completed", - usage=None, - ) +_CHATGPT_RESPONSES_URL: Final = "https://chatgpt.com/backend-api/codex/responses" -class _CompletedStream: - def __init__(self, response): - self.completed_response = MagicMock(response=response) - self._hidden_params = {} - - def __iter__(self): - return iter(()) - - def __aiter__(self): - return self - - async def __anext__(self): - raise StopAsyncIteration +def _reasoning_item(item_id: str) -> dict[str, object]: + return {"type": "reasoning", "id": item_id, "summary": [], "encrypted_content": f"gAAAA_{item_id}"} -def _messages_client_kwargs(call_type: str): - messages = [{"role": "user", "content": "Weather in Oslo and Lyon?"}] +def _weather_call(index: int, city: str) -> dict[str, object]: return { - "model": "gpt-5.4-nano", - "custom_llm_provider": "azure", - "messages": messages, - "optional_params": {}, - "litellm_params": {}, - "headers": {}, - "model_response": ModelResponse(), - "logging_obj": LiteLLMLogging( - litellm_call_id="test-call", - call_type=call_type, - model="gpt-5.4-nano", - messages=messages, - function_id="fn-id", - stream=False, - start_time=datetime.now(), - ), + "type": "function_call", + "id": f"fc_{index}", + "status": "completed", + "arguments": json.dumps({"city": city}), + "call_id": f"call_{index}", + "name": "get_weather", } -@pytest.mark.parametrize("is_async", [True, False], ids=["acompletion", "completion"]) -@pytest.mark.parametrize("arrives_streamed", [False, True], ids=["response", "collected_stream"]) -@pytest.mark.parametrize( - ("call_type", "kept_ids"), - [ - pytest.param("anthropic_messages", ["rs_first", "rs_second"], id="messages_client"), - pytest.param("acompletion", ["rs_second"], id="chat_client"), - ], -) +def _two_reasoning_tool_turn_stream() -> httpx.Response: + turn: Final = { + "id": "resp_two_reasoning", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6-sol", + "output": [ + _reasoning_item("rs_first"), + _weather_call(1, "Oslo"), + _reasoning_item("rs_second"), + _weather_call(2, "Lyon"), + ], + "usage": {"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, + } + body: Final = f"data: {json.dumps({'type': 'response.completed', 'response': turn})}\n\ndata: [DONE]\n\n" + return httpx.Response(200, headers={"content-type": "text/event-stream"}, text=body) + + +@pytest.fixture +def _chatgpt_backend(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter) -> None: + auth: Final = {"access_token": "chatgpt-test-token", "expires_at": time.time() + 3600, "account_id": "acct_test"} + (tmp_path / "auth.json").write_text(json.dumps(auth)) + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(_CHATGPT_RESPONSES_URL).mock(side_effect=lambda _request: _two_reasoning_tool_turn_stream()) + + +_WEATHER_QUESTION: Final = "Weather in Oslo and Lyon?" +_WEATHER_TOOL_INPUT: Final = {"type": "object", "properties": {"city": {"type": "string"}}} +_EVERY_REASONING_SIGNATURE: Final = [ + encrypted_reasoning_signature("gAAAA_rs_first"), + encrypted_reasoning_signature("gAAAA_rs_second"), +] + + +def _thinking_signatures(reply: object) -> list[object]: + assert isinstance(reply, dict) + return [block.get("data") for block in reply["content"] if block["type"] == "redacted_thinking"] + + @pytest.mark.asyncio -async def test_non_streamed_bridge_hands_messages_clients_every_reasoning_item( - is_async: bool, arrives_streamed: bool, call_type: str, kept_ids: list[str] -): - turn = _two_reasoning_tool_turn() - upstream = _CompletedStream(turn) if arrives_streamed else turn - bridge = ResponsesToCompletionBridgeHandler() - kwargs = _messages_client_kwargs(call_type) +@pytest.mark.usefixtures("_chatgpt_backend") +async def test_async_messages_client_gets_every_reasoning_item_of_a_non_streamed_turn() -> None: + reply: Final = await litellm.anthropic.messages.acreate( + model="chatgpt/gpt-5.6-sol", + max_tokens=64, + messages=[{"role": "user", "content": _WEATHER_QUESTION}], + tools=[{"name": "get_weather", "input_schema": _WEATHER_TOOL_INPUT}], + ) - if is_async: - with patch("litellm.aresponses", new=AsyncMock(return_value=upstream)): - result = await bridge.acompletion(**kwargs) - else: - with patch("litellm.responses", return_value=upstream): - result = bridge.completion(**kwargs) + assert _thinking_signatures(reply) == _EVERY_REASONING_SIGNATURE - assert [item["id"] for item in result.choices[0].message.reasoning_items] == kept_ids + +_CHAT_WEATHER_TOOL: Final = {"type": "function", "function": {"name": "get_weather", "parameters": _WEATHER_TOOL_INPUT}} + + +def _reasoning_item_ids(response: object) -> list[object]: + assert isinstance(response, ModelResponse) + return [item["id"] for item in response.choices[0].message.reasoning_items] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chatgpt_backend") +async def test_async_chat_client_keeps_only_the_last_reasoning_item_of_a_turn() -> None: + response: Final = await litellm.acompletion( + model="chatgpt/gpt-5.6-sol", + messages=[{"role": "user", "content": _WEATHER_QUESTION}], + tools=[_CHAT_WEATHER_TOOL], + ) + + assert _reasoning_item_ids(response) == ["rs_second"] + + +@pytest.mark.usefixtures("_chatgpt_backend") +def test_sync_chat_client_keeps_only_the_last_reasoning_item_of_a_turn() -> None: + response: Final = litellm.completion( + model="chatgpt/gpt-5.6-sol", + messages=[{"role": "user", "content": _WEATHER_QUESTION}], + tools=[_CHAT_WEATHER_TOOL], + ) + + assert _reasoning_item_ids(response) == ["rs_second"] diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 1223185a05c..0463d37c4a8 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -3,6 +3,7 @@ import functools import json import logging from collections.abc import Mapping, Sequence +from contextlib import nullcontext from types import MappingProxyType from typing import Any, Final, cast @@ -11,9 +12,7 @@ import pytest import respx import litellm - - - +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_PLACEHOLDER, encrypted_reasoning_signature, @@ -5585,6 +5584,85 @@ def test_sync_messages_retry_once_without_reasoning_the_backend_cannot_decrypt( assert [block["text"] for block in response["content"] if block["type"] == "text"] == ["Oslo"] +class _FailureLog(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.rejections: tuple[str, ...] = () + + def _record(self, kwargs) -> None: + error: Final = str(kwargs.get("exception")) + self.rejections = ( + *self.rejections, + next((phrase for phrase in _UPSTREAM_REJECTION_PHRASES if phrase in error), error), + ) + + def log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + self._record(kwargs) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + self._record(kwargs) + + +_UPSTREAM_REJECTION_PHRASES: Final = ("could not be verified", "max_output_tokens is too small") +_OUTCOMES_OF_THE_TURN_SENT_AGAIN: Final = ( + pytest.param(httpx.Response(200, json=_RETRY_REPLY), (), id="served"), + pytest.param(_UNRELATED_REJECTION, ("max_output_tokens is too small",), id="rejected"), +) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +@pytest.mark.parametrize(("sent_again", "logged_rejections"), _OUTCOMES_OF_THE_TURN_SENT_AGAIN) +async def test_async_messages_log_only_the_outcome_of_the_turn_sent_again( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + sent_again: httpx.Response, + logged_rejections: tuple[str, ...], +) -> None: + failure_log: Final = _FailureLog() + monkeypatch.setattr(litellm, "callbacks", [failure_log]) + respx_mock.post(url__regex=_RETRY_URL).mock(side_effect=[_RETRY_REJECTION, sent_again]) + + with pytest.raises(litellm.BadRequestError) if logged_rejections else nullcontext(): + await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + ) + + assert failure_log.rejections == logged_rejections + + +@pytest.mark.usefixtures("_chat_bridge_httpx") +@pytest.mark.parametrize(("sent_again", "logged_rejections"), _OUTCOMES_OF_THE_TURN_SENT_AGAIN) +def test_sync_messages_log_only_the_outcome_of_the_turn_sent_again( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + sent_again: httpx.Response, + logged_rejections: tuple[str, ...], +) -> None: + failure_log: Final = _FailureLog() + monkeypatch.setattr(litellm, "callbacks", [failure_log]) + respx_mock.post(url__regex=_RETRY_URL).mock(side_effect=[_RETRY_REJECTION, sent_again]) + + with pytest.raises(litellm.BadRequestError) if logged_rejections else nullcontext(): + litellm.anthropic.messages.create( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + ) + + assert failure_log.rejections == logged_rejections + + @pytest.mark.asyncio @pytest.mark.usefixtures("_chat_bridge_httpx") @pytest.mark.parametrize( diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py index 042f99be001..c11023b256a 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py @@ -1,8 +1,10 @@ +import asyncio import datetime import json import os import sys from collections.abc import AsyncIterator, Mapping +from contextlib import nullcontext from typing import Final, cast from unittest.mock import AsyncMock, patch @@ -13,6 +15,8 @@ import respx sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))) import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, @@ -795,6 +799,85 @@ def test_sync_openai_turn_is_sent_again_once_without_replayed_reasoning_the_back assert cast(Mapping[str, object], response)["content"] == [{"type": "text", "text": "done"}] +class _FailureLog(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.rejections: tuple[str, ...] = () + + def _record(self, kwargs) -> None: + error: Final = str(kwargs.get("exception")) + self.rejections = ( + *self.rejections, + next((phrase for phrase in ("encrypted content", "organization") if phrase in error), error), + ) + + def log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + self._record(kwargs) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + self._record(kwargs) + + +_OUTCOMES_OF_THE_TURN_SENT_AGAIN: Final = ( + pytest.param(httpx.Response(200, json=TEXT_RESPONSE), (), id="served"), + pytest.param(httpx.Response(400, json=UNVERIFIED_ORGANIZATION), ("organization",), id="rejected"), +) + + +def _rejected_then(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sent_again: httpx.Response) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-replay-test") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[httpx.Response(400, json=REJECTED_REPLAYED_REASONING), sent_again] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("sent_again", "logged_rejections"), _OUTCOMES_OF_THE_TURN_SENT_AGAIN) +async def test_openai_turn_logs_only_the_outcome_of_the_turn_sent_again( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + sent_again: httpx.Response, + logged_rejections: tuple[str, ...], +) -> None: + failure_log: Final = _FailureLog() + monkeypatch.setattr(litellm, "callbacks", [failure_log]) + _rejected_then(respx_mock, monkeypatch, sent_again) + + with pytest.raises(litellm.BadRequestError) if logged_rejections else nullcontext(): + await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + ) + + assert failure_log.rejections == logged_rejections + + +@pytest.mark.parametrize(("sent_again", "logged_rejections"), _OUTCOMES_OF_THE_TURN_SENT_AGAIN) +def test_sync_openai_turn_logs_only_the_outcome_of_the_turn_sent_again( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + sent_again: httpx.Response, + logged_rejections: tuple[str, ...], +) -> None: + failure_log: Final = _FailureLog() + monkeypatch.setattr(litellm, "callbacks", [failure_log]) + _rejected_then(respx_mock, monkeypatch, sent_again) + + with pytest.raises(litellm.BadRequestError) if logged_rejections else nullcontext(): + litellm.anthropic.messages.create( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + ) + + assert failure_log.rejections == logged_rejections + + @pytest.mark.asyncio @pytest.mark.parametrize( "replayed_reasoning, rejection", @@ -821,3 +904,59 @@ async def test_openai_turn_is_not_sent_again_for_other_rejections( ) assert upstream.call_count == 1 + + +class _SuccessLog(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged_reasoning_item_ids: tuple[tuple[object, ...], ...] = () + self.logged: Final = asyncio.Event() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + logged: Final = kwargs["standard_logging_object"]["response"] + items: Final = logged["choices"][0]["message"].get("reasoning_items") or () + self.logged_reasoning_item_ids = (*self.logged_reasoning_item_ids, tuple(item["id"] for item in items)) + self.logged.set() + + +def _reasoning_before_call(index: int) -> tuple[dict[str, object], dict[str, object]]: + return ( + {"type": "reasoning", "id": f"rs_{index}", "summary": [], "encrypted_content": f"gAAA_rs_{index}"}, + { + "type": "function_call", + "id": f"fc_{index}", + "status": "completed", + "arguments": "{}", + "call_id": f"call_{index}", + "name": "Bash", + }, + ) + + +@pytest.mark.asyncio +async def test_openai_turn_logs_every_reasoning_item_the_messages_client_was_served( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + success_log: Final = _SuccessLog() + monkeypatch.setattr(litellm, "callbacks", [success_log]) + monkeypatch.setenv("OPENAI_API_KEY", "sk-replay-test") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response( + 200, + json={**TEXT_RESPONSE, "output": [*_reasoning_before_call(1), *_reasoning_before_call(2)]}, + ) + ) + + reply: Final = await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=[{"role": "user", "content": "List the logs twice."}], + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + ) + GLOBAL_LOGGING_WORKER.start() + await asyncio.wait_for(success_log.logged.wait(), timeout=30.0) + + assert [block["type"] for block in reply["content"]].count("redacted_thinking") == 2 + assert success_log.logged_reasoning_item_ids == (("rs_1", "rs_2"),) From 5a6475d112037d7968e7b52f31016760f96f66d6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 08:46:14 -0800 Subject: [PATCH 28/41] test(anthropic): type the parameters of the failure-recording test loggers Backport of 3d53f0047d1 (#49076) to lts/1.108.x. --- ...xperimental_pass_through_adapters_transformation.py | 10 +++++++--- .../test_responses_adapters_handler.py | 10 +++++++--- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 0463d37c4a8..6222aa7aab5 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -5589,17 +5589,21 @@ class _FailureLog(CustomLogger): super().__init__() self.rejections: tuple[str, ...] = () - def _record(self, kwargs) -> None: + def _record(self, kwargs: Mapping[str, object]) -> None: error: Final = str(kwargs.get("exception")) self.rejections = ( *self.rejections, next((phrase for phrase in _UPSTREAM_REJECTION_PHRASES if phrase in error), error), ) - def log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: self._record(kwargs) - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: self._record(kwargs) diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py index c11023b256a..320fb6123bf 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py @@ -804,17 +804,21 @@ class _FailureLog(CustomLogger): super().__init__() self.rejections: tuple[str, ...] = () - def _record(self, kwargs) -> None: + def _record(self, kwargs: Mapping[str, object]) -> None: error: Final = str(kwargs.get("exception")) self.rejections = ( *self.rejections, next((phrase for phrase in ("encrypted content", "organization") if phrase in error), error), ) - def log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: self._record(kwargs) - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: self._record(kwargs) From c72cb031d8fdb33bb37d30e281847f4dc856e563 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 08:55:37 -0800 Subject: [PATCH 29/41] test(anthropic): type the success logger and use a fixed token expiry in the bridge tests --- ...tras_litellm_responses_transformation_handler.py | 6 ++++-- .../test_responses_adapters_handler.py | 13 +++++++++---- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py index af0a0dbe9b0..ab31cee5e07 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py @@ -1,5 +1,4 @@ import json -import time from datetime import datetime from pathlib import Path from typing import Final @@ -381,9 +380,12 @@ def _two_reasoning_tool_turn_stream() -> httpx.Response: return httpx.Response(200, headers={"content-type": "text/event-stream"}, text=body) +_FAR_FUTURE_EXPIRY: Final = 4_102_444_800 + + @pytest.fixture def _chatgpt_backend(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter) -> None: - auth: Final = {"access_token": "chatgpt-test-token", "expires_at": time.time() + 3600, "account_id": "acct_test"} + auth: Final = {"access_token": "chatgpt-test-token", "expires_at": _FAR_FUTURE_EXPIRY, "account_id": "acct_test"} (tmp_path / "auth.json").write_text(json.dumps(auth)) monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py index 320fb6123bf..f21eb9eb6ac 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py @@ -3,7 +3,7 @@ import datetime import json import os import sys -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Mapping, Sequence from contextlib import nullcontext from typing import Final, cast from unittest.mock import AsyncMock, patch @@ -28,6 +28,7 @@ from litellm.llms.anthropic.pass_through.responses_adapters.handler import ( ) from litellm.router import Router from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS, converted_stream_requested +from litellm.types.utils import StandardLoggingPayload MESSAGES = [{"role": "user", "content": "hello"}] CLAUDE_CODE_USER_ID = json.dumps({"device_id": "d" * 64, "account_uuid": "", "session_id": "session-abc"}) @@ -916,9 +917,13 @@ class _SuccessLog(CustomLogger): self.logged_reasoning_item_ids: tuple[tuple[object, ...], ...] = () self.logged: Final = asyncio.Event() - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: - logged: Final = kwargs["standard_logging_object"]["response"] - items: Final = logged["choices"][0]["message"].get("reasoning_items") or () + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + payload: Final = cast(StandardLoggingPayload, kwargs["standard_logging_object"]) + logged: Final = cast(Mapping[str, Sequence[Mapping[str, Mapping[str, object]]]], payload["response"]) + message: Final = logged["choices"][0]["message"] + items: Final = cast(Sequence[Mapping[str, object]], message.get("reasoning_items") or ()) self.logged_reasoning_item_ids = (*self.logged_reasoning_item_ids, tuple(item["id"] for item in items)) self.logged.set() From 433a76f754a4bfc4a804f40a3721be79df07ff33 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 09:28:59 -0800 Subject: [PATCH 30/41] fix(anthropic): skip the SDK num_retries retry for an attempt that is sent again --- litellm/_internal_context.py | 3 +- litellm/utils.py | 14 ++++-- ...al_pass_through_adapters_transformation.py | 45 +++++++++++++++++++ .../test_responses_adapters_handler.py | 33 ++++++++++++++ 4 files changed, 90 insertions(+), 5 deletions(-) diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index 1b07c017a3c..c4527e1d099 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -118,7 +118,8 @@ def current_service_caller() -> str | None: @contextmanager def resent_on(rejection: Callable[[Exception], bool]) -> Generator[None]: """An attempt its caller sends again when it fails with ``rejection``. That failure is not the request's - outcome, so nothing logs it and its budget reservation and rate-limit slot stay with the attempt sent again.""" + outcome, so it is neither logged as the request's failure nor retried, and its budget reservation and rate-limit + slot stay with the attempt sent again.""" token: Final = _resent_on.set(rejection) try: yield diff --git a/litellm/utils.py b/litellm/utils.py index 3213e343c4f..4e1c6a91130 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1985,7 +1985,7 @@ def client(original_function): is_completion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=False) if ( - num_retries and not is_completion_litellm_router_call + num_retries and not is_completion_litellm_router_call and not is_resent_failure(e) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying if ( isinstance(e, openai.APIError) @@ -2019,7 +2019,7 @@ def client(original_function): is_responses_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=False) if ( - num_retries and not is_responses_litellm_router_call + num_retries and not is_responses_litellm_router_call and not is_resent_failure(e) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying if ( isinstance(e, openai.APIError) @@ -2300,7 +2300,10 @@ def client(original_function): is_acompletion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True) if ( - num_retries and not is_acompletion_litellm_router_call and not isinstance(e, ImportError) + num_retries + and not is_acompletion_litellm_router_call + and not isinstance(e, ImportError) + and not is_resent_failure(e) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying try: litellm.num_retries = None # set retries to None to prevent infinite loops @@ -2331,7 +2334,10 @@ def client(original_function): is_aresponses_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True) if ( - num_retries and not is_aresponses_litellm_router_call and not isinstance(e, ImportError) + num_retries + and not is_aresponses_litellm_router_call + and not isinstance(e, ImportError) + and not is_resent_failure(e) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying try: litellm.num_retries = None # set retries to None to prevent infinite loops diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 6222aa7aab5..1683b4a41e7 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -5584,6 +5584,51 @@ def test_sync_messages_retry_once_without_reasoning_the_backend_cannot_decrypt( assert [block["text"] for block in response["content"] if block["type"] == "text"] == ["Oslo"] +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_async_messages_with_num_retries_send_the_rejected_reasoning_only_once( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=_RETRY_URL).mock( + side_effect=[_RETRY_REJECTION, httpx.Response(200, json=_RETRY_REPLY)] + ) + + await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + num_retries=2, + ) + + assert [len(_sent_reasoning_items(call.request)) for call in route.calls] == [1, 0] + + +@pytest.mark.usefixtures("_chat_bridge_httpx") +def test_sync_messages_with_num_retries_send_the_rejected_reasoning_only_once( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=_RETRY_URL).mock( + side_effect=[_RETRY_REJECTION, httpx.Response(200, json=_RETRY_REPLY)] + ) + + litellm.anthropic.messages.create( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + num_retries=2, + ) + + assert [len(_sent_reasoning_items(call.request)) for call in route.calls] == [1, 0] + + class _FailureLog(CustomLogger): def __init__(self) -> None: super().__init__() diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py index f21eb9eb6ac..19de6ccd634 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py @@ -800,6 +800,39 @@ def test_sync_openai_turn_is_sent_again_once_without_replayed_reasoning_the_back assert cast(Mapping[str, object], response)["content"] == [{"type": "text", "text": "done"}] +@pytest.mark.asyncio +async def test_openai_turn_with_num_retries_sends_the_rejected_reasoning_only_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + upstream: Final = _upstream(respx_mock, monkeypatch, httpx.Response(400, json=REJECTED_REPLAYED_REASONING), False) + + await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + num_retries=2, + ) + + assert _replayed_reasoning_counts(upstream) == [1, 0] + + +def test_sync_openai_turn_with_num_retries_sends_the_rejected_reasoning_only_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + upstream: Final = _upstream(respx_mock, monkeypatch, httpx.Response(400, json=REJECTED_REPLAYED_REASONING), False) + + litellm.anthropic.messages.create( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + num_retries=2, + ) + + assert _replayed_reasoning_counts(upstream) == [1, 0] + + class _FailureLog(CustomLogger): def __init__(self) -> None: super().__init__() From b277d687b5e6f596b8cd26a5286ab2bc4c2f0034 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 10:22:13 -0800 Subject: [PATCH 31/41] fix(anthropic): stop SDK retries from resending a rejected block or clearing the global num_retries A transient 429 or 500 on a /v1/messages turn that carries a replayed encrypted block was retried by an SDK call's num_retries loop, which sent the block again until every retry was rejected. On async calls the loop then raised the first transient error, so the turn failed instead of being sent again without the block. The retry loop now stops at that rejection and passes it on, so the turn is sent again once without encrypted blocks. A rejected first try on a sync call also cleared the process-wide litellm.num_retries, so later calls in the process were no longer retried. That setting now stays as configured. --- litellm/main.py | 33 ++++++++-- litellm/utils.py | 16 +++-- ...al_pass_through_adapters_transformation.py | 59 ++++++++++++++++- .../test_responses_adapters_handler.py | 63 ++++++++++++++++++- 4 files changed, 155 insertions(+), 16 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 42d5ecc8f20..8626d64ebd7 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -27,6 +27,7 @@ from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args +from litellm._internal_context import is_resent_failure from litellm._logging import redact_string from litellm._uuid import uuid @@ -6242,6 +6243,10 @@ def completion( ) +def _is_not_resent_failure(error: BaseException) -> bool: + return not (isinstance(error, Exception) and is_resent_failure(error)) + + def completion_with_retries(*args, **kwargs): """ Executes a litellm.completion() with 3 retries @@ -6261,10 +6266,15 @@ def completion_with_retries(*args, **kwargs): retryer = tenacity.Retrying( wait=tenacity.wait_exponential(multiplier=1, max=10), stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), reraise=True, ) else: - retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + retryer = tenacity.Retrying( + stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), + reraise=True, + ) return retryer(original_function, *args, **kwargs) @@ -6287,10 +6297,15 @@ async def acompletion_with_retries(*args, **kwargs): retryer = tenacity.AsyncRetrying( wait=tenacity.wait_exponential(multiplier=1, max=10), stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), reraise=True, ) else: - retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + retryer = tenacity.AsyncRetrying( + stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), + reraise=True, + ) return await retryer(original_function, *args, **kwargs) @@ -6315,10 +6330,15 @@ def responses_with_retries(*args, **kwargs): retryer = tenacity.Retrying( wait=tenacity.wait_exponential(multiplier=1, max=10), stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), reraise=True, ) else: - retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + retryer = tenacity.Retrying( + stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), + reraise=True, + ) return retryer(original_function, *args, **kwargs) @@ -6342,10 +6362,15 @@ async def aresponses_with_retries(*args, **kwargs): retryer = tenacity.AsyncRetrying( wait=tenacity.wait_exponential(multiplier=1, max=10), stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), reraise=True, ) else: - retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + retryer = tenacity.AsyncRetrying( + stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), + reraise=True, + ) return await retryer(original_function, *args, **kwargs) diff --git a/litellm/utils.py b/litellm/utils.py index 4e1c6a91130..c2acb89b6e6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1980,7 +1980,8 @@ def client(original_function): retry_policy=kwargs.get("retry_policy"), ) kwargs["retry_policy"] = reset_retry_policy() # prevent infinite loops - litellm.num_retries = None # set retries to None to prevent infinite loops + if not is_resent_failure(e): + litellm.num_retries = None # set retries to None to prevent infinite loops context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict", {}) is_completion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=False) @@ -2015,7 +2016,8 @@ def client(original_function): retry_policy=kwargs.get("retry_policy"), ) kwargs["retry_policy"] = reset_retry_policy() # prevent infinite loops - litellm.num_retries = None # set retries to None to prevent infinite loops + if not is_resent_failure(e): + litellm.num_retries = None # set retries to None to prevent infinite loops is_responses_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=False) if ( @@ -2314,8 +2316,9 @@ def client(original_function): elif isinstance(e, openai.APIError): # generic api error kwargs["retry_strategy"] = "constant_retry" result = await litellm.acompletion_with_retries(*args, **kwargs) - except Exception: - pass + except Exception as retry_error: + if is_resent_failure(retry_error): + raise else: return result elif ( @@ -2348,8 +2351,9 @@ def client(original_function): elif isinstance(e, openai.APIError): # generic api error kwargs["retry_strategy"] = "constant_retry" result = await litellm.aresponses_with_retries(*args, **kwargs) - except Exception: - pass + except Exception as retry_error: + if is_resent_failure(retry_error): + raise else: return result diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 1683b4a41e7..f02e7762deb 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -5608,9 +5608,10 @@ async def test_async_messages_with_num_retries_send_the_rejected_reasoning_only_ @pytest.mark.usefixtures("_chat_bridge_httpx") -def test_sync_messages_with_num_retries_send_the_rejected_reasoning_only_once( - respx_mock: respx.MockRouter, +def test_sync_messages_with_global_num_retries_send_the_rejected_reasoning_once_and_keep_the_setting( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: + monkeypatch.setattr(litellm, "num_retries", 2) route: Final = respx_mock.post(url__regex=_RETRY_URL).mock( side_effect=[_RETRY_REJECTION, httpx.Response(200, json=_RETRY_REPLY)] ) @@ -5623,10 +5624,62 @@ def test_sync_messages_with_num_retries_send_the_rejected_reasoning_only_once( max_tokens=64, tools=_RETRY_TOOLS, messages=_history_with_foreign_reasoning(), - num_retries=2, ) assert [len(_sent_reasoning_items(call.request)) for call in route.calls] == [1, 0] + assert litellm.num_retries == 2 + + +_SERVER_ERROR: Final = httpx.Response( + 500, json={"error": {"message": "The server had an error while processing your request.", "type": "server_error"}} +) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_async_messages_retry_a_server_error_with_the_reasoning_and_drop_it_only_once_rejected( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=_RETRY_URL).mock( + side_effect=[_SERVER_ERROR, _RETRY_REJECTION, httpx.Response(200, json=_RETRY_REPLY)] + ) + + response: Final = await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + num_retries=2, + ) + + assert [len(_sent_reasoning_items(call.request)) for call in route.calls] == [1, 1, 0] + assert [block["text"] for block in response["content"] if block["type"] == "text"] == ["Oslo"] + + +@pytest.mark.usefixtures("_chat_bridge_httpx") +def test_sync_messages_retry_a_server_error_with_the_reasoning_and_drop_it_only_once_rejected( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=_RETRY_URL).mock( + side_effect=[_SERVER_ERROR, _RETRY_REJECTION, httpx.Response(200, json=_RETRY_REPLY)] + ) + + response: Final = litellm.anthropic.messages.create( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + num_retries=2, + ) + + assert [len(_sent_reasoning_items(call.request)) for call in route.calls] == [1, 1, 0] + assert [block["text"] for block in response["content"] if block["type"] == "text"] == ["Oslo"] class _FailureLog(CustomLogger): diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py index 19de6ccd634..526ead24d26 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py @@ -747,6 +747,12 @@ def _replayed_reasoning_counts(upstream: respx.Route) -> list[int]: def _upstream(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, first: httpx.Response, stream: bool): + return _upstream_after(respx_mock, monkeypatch, (first,), stream) + + +def _upstream_after( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, failures: tuple[httpx.Response, ...], stream: bool +) -> respx.Route: monkeypatch.setenv("OPENAI_API_KEY", "sk-replay-test") monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") litellm.in_memory_llm_clients_cache.flush_cache() @@ -755,7 +761,7 @@ def _upstream(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, fir if stream else httpx.Response(200, json=TEXT_RESPONSE) ) - return respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=[first, served]) + return respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=[*failures, served]) @pytest.mark.asyncio @@ -817,9 +823,10 @@ async def test_openai_turn_with_num_retries_sends_the_rejected_reasoning_only_on assert _replayed_reasoning_counts(upstream) == [1, 0] -def test_sync_openai_turn_with_num_retries_sends_the_rejected_reasoning_only_once( +def test_sync_openai_turn_with_global_num_retries_sends_the_rejected_reasoning_once_and_keeps_the_setting( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: + monkeypatch.setattr(litellm, "num_retries", 2) upstream: Final = _upstream(respx_mock, monkeypatch, httpx.Response(400, json=REJECTED_REPLAYED_REASONING), False) litellm.anthropic.messages.create( @@ -827,10 +834,60 @@ def test_sync_openai_turn_with_num_retries_sends_the_rejected_reasoning_only_onc messages=_history(replayed_reasoning=True), model="openai/gpt-5.6-luna", tools=[dict(BASH_TOOL)], - num_retries=2, ) assert _replayed_reasoning_counts(upstream) == [1, 0] + assert litellm.num_retries == 2 + + +_SERVER_ERROR: Final = { + "error": {"message": "The server had an error while processing your request.", "type": "server_error"} +} + + +@pytest.mark.asyncio +async def test_openai_turn_retries_a_server_error_with_the_reasoning_and_drops_it_only_once_rejected( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + upstream: Final = _upstream_after( + respx_mock, + monkeypatch, + (httpx.Response(500, json=_SERVER_ERROR), httpx.Response(400, json=REJECTED_REPLAYED_REASONING)), + False, + ) + + response: Final = await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + num_retries=2, + ) + + assert _replayed_reasoning_counts(upstream) == [1, 1, 0] + assert cast(Mapping[str, object], response)["content"] == [{"type": "text", "text": "done"}] + + +def test_sync_openai_turn_retries_a_server_error_with_the_reasoning_and_drops_it_only_once_rejected( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + upstream: Final = _upstream_after( + respx_mock, + monkeypatch, + (httpx.Response(500, json=_SERVER_ERROR), httpx.Response(400, json=REJECTED_REPLAYED_REASONING)), + False, + ) + + response: Final = litellm.anthropic.messages.create( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + num_retries=2, + ) + + assert _replayed_reasoning_counts(upstream) == [1, 1, 0] + assert cast(Mapping[str, object], response)["content"] == [{"type": "text", "text": "done"}] class _FailureLog(CustomLogger): From 3eda243ea828153f4a539c477b7a269a7f55aa08 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 7 Nov 2026 10:31:40 -0800 Subject: [PATCH 32/41] chore: bump litellm 1.108.2 -> 1.108.3 litellm 1.108.2 already shipped, so the next release off this line is 1.108.3 --- pyproject.toml | 4 ++-- uv.lock | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index fe156695894..84200cb91a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.108.2" +version = "1.108.3" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.15" @@ -396,7 +396,7 @@ litellm-enterprise = { workspace = true } members = ["enterprise", "litellm-proxy-extras"] [tool.commitizen] -version = "1.108.2" +version = "1.108.3" version_files = [ "pyproject.toml:^version", ] diff --git a/uv.lock b/uv.lock index cc78423f359..0df2668e679 100644 --- a/uv.lock +++ b/uv.lock @@ -4594,7 +4594,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.108.2" +version = "1.108.3" source = { editable = "." } dependencies = [ { name = "aiohttp" }, From edb48c1d09222b6d18899c5ea0f59b43bebb7ae4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 10:45:18 -0800 Subject: [PATCH 33/41] fix(anthropic): stop a cancelled SDK retry from starting another call The resent-failure retry predicate returned true for BaseException subclasses such as asyncio.CancelledError, so a cancelled acompletion_with_retries or aresponses_with_retries call could start another provider call. Retry only Exception subclasses again, as tenacity's default did, while still skipping a resent failure. --- litellm/main.py | 2 +- tests/unit/test_main.py | 11 +++++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/litellm/main.py b/litellm/main.py index 8626d64ebd7..586f42311b4 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6244,7 +6244,7 @@ def completion( def _is_not_resent_failure(error: BaseException) -> bool: - return not (isinstance(error, Exception) and is_resent_failure(error)) + return isinstance(error, Exception) and not is_resent_failure(error) def completion_with_retries(*args, **kwargs): diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 6bb69def307..150f45e009d 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -7936,6 +7936,17 @@ async def test_responses_with_retries(sync_mode): assert mock_responses.call_args.kwargs["max_retries"] == 0 +@pytest.mark.parametrize("retry_function", [acompletion_with_retries, aresponses_with_retries]) +@pytest.mark.asyncio +async def test_async_retry_helpers_stop_when_the_call_is_cancelled(retry_function): + upstream: Final = AsyncMock(side_effect=[asyncio.CancelledError(), MagicMock()]) + + with pytest.raises(asyncio.CancelledError): + await retry_function(model="gpt-4o", num_retries=3, original_function=upstream) + + assert upstream.await_count == 1 + + def test_azure_embedding_exceptions(): with pytest.raises(Exception, match="Mock error") as exc_info: litellm.embedding( From d5ad71dd6584d9a942edf45a7fe0f346d3455440 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 10:54:21 -0800 Subject: [PATCH 34/41] test(anthropic): type the retry helper parameter in the cancellation test --- tests/unit/test_main.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 150f45e009d..c832da5dbfc 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -8,7 +8,7 @@ import logging import os import threading import urllib.parse -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from datetime import datetime @@ -7938,7 +7938,9 @@ async def test_responses_with_retries(sync_mode): @pytest.mark.parametrize("retry_function", [acompletion_with_retries, aresponses_with_retries]) @pytest.mark.asyncio -async def test_async_retry_helpers_stop_when_the_call_is_cancelled(retry_function): +async def test_async_retry_helpers_stop_when_the_call_is_cancelled( + retry_function: Callable[..., Awaitable[object]], +) -> None: upstream: Final = AsyncMock(side_effect=[asyncio.CancelledError(), MagicMock()]) with pytest.raises(asyncio.CancelledError): From 52f49b8e8628e5360aef3710848b7780a6d594cb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 7 Nov 2026 14:34:39 -0800 Subject: [PATCH 35/41] test(integration): backport main's Gateway.rows, Scenario.model name and MCP ack fix to lts/1.108.x Backports test-only changes from main so the line's integration signal stops reporting stale helpers as failures. Gateway.database_url and Gateway.rows come from #47213 (fc27dfdc7d), and the name keyword on Scenario.model comes from #48917 (ebf84f8649). Only those hunks of tests/integration/_support/client.py are taken. The rest of #47213 gives each owned proxy its own database and Redis, which none of the line's tests need, and the workers field from a later main change is left out. The MCP progress and cancellation test wraps its upstream in _acknowledge_notifications_after_dispatch, from #49967 (b5952f243a). The helper is already on the line. --- tests/integration/_support/client.py | 14 ++++++++++---- tests/integration/mcp/test_mcp_transports.py | 6 ++++-- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index 053318c2451..fe1db22987d 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -95,6 +95,10 @@ class Gateway: client: httpx.Client key: str upstream_url: str + database_url: str | None = None + + def rows(self, query: str, parameters: tuple[str, ...]) -> list[dict[str, JsonValue]]: + return read_rows(query, parameters, database_url=self.database_url) def request( self, @@ -267,12 +271,14 @@ class Scenario: ) assert read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == [] - def model(self, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue) -> str: - name: Final = f"integration-{uuid.uuid4().hex}" + def model( + self, *, name: str | None = None, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue + ) -> str: + model_name: Final = name if name is not None else f"integration-{uuid.uuid4().hex}" created: Final = self.gateway.post( "/model/new", { - "model_name": name, + "model_name": model_name, "litellm_params": { "model": "openai/gpt-4o-mini", "api_key": "integration-provider-key", @@ -284,7 +290,7 @@ class Scenario: ) identity: Final = string_value(object_value(created["model_info"])["id"]) self.cleanups.callback(self.delete_model, identity) - return name + return model_name @contextmanager diff --git a/tests/integration/mcp/test_mcp_transports.py b/tests/integration/mcp/test_mcp_transports.py index ef240cba570..be9e016e947 100644 --- a/tests/integration/mcp/test_mcp_transports.py +++ b/tests/integration/mcp/test_mcp_transports.py @@ -678,7 +678,7 @@ def test_sdk_progress_and_cancellation_reach_upstream_without_losing_session( from starlette.routing import Mount from integration._support.asgi import asgi_server - from integration._support.mcp import _drain_sse_streams + from integration._support.mcp import _acknowledge_notifications_after_dispatch, _drain_sse_streams from integration._support.process import owned_proxy upstream: Final = MCPServer("progress-cancellation-peer") @@ -770,7 +770,9 @@ def test_sdk_progress_and_cancellation_reach_upstream_without_losing_session( watchers = tuple(task for task in asyncio.all_tasks() if "_shutdown_watcher" in repr(task.get_coro())) await asyncio.wait_for(asyncio.gather(*watchers), 2) - hosted: Final = Starlette(routes=[Mount("/", app=app)], lifespan=lifespan) + hosted: Final = Starlette( + routes=[Mount("/", app=_acknowledge_notifications_after_dispatch(app))], lifespan=lifespan + ) with asgi_server(hosted) as upstream_url, httpx.Client() as client: config: Final = tmp_path / "proxy.yaml" config.write_text( From f4ba1bfafc682de903b04ce0cf2e81566977a2db Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 7 Nov 2026 16:45:18 -0800 Subject: [PATCH 36/41] fix(bedrock)!: backport #47429 deployment guardrailConfig on /bedrock pass-through routes to lts/1.108.x The router-model pass-through dropped a Bedrock deployment's guardrailConfig, so converse, converse-stream, invoke and invoke-with-response-stream reached Bedrock unguarded while /v1/chat/completions on the same deployment was blocked. The signer now sends it the way Bedrock reads it, in the body for Converse and as the X-Amzn-Bedrock-Guardrail* headers for Invoke, replacing any guardrail the caller sent Adapted to the line, which does not have #47040. The guardrail headers are merged through merge_bedrock_invoke_headers without the request metadata header, and the unit tests drop the request metadata assertions and the identity argument that only #47040 needs --- .../bedrock/passthrough/transformation.py | 46 +- litellm/llms/bedrock/request_metadata.py | 11 +- litellm/passthrough/main.py | 3 +- ...drock_passthrough_guardrail_config_wire.py | 531 ++++++++++++++++++ ...test_bedrock_passthrough_transformation.py | 195 ++++++- 5 files changed, 776 insertions(+), 10 deletions(-) create mode 100644 tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index d7d6341fe7a..49eea6f5daf 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -13,6 +13,10 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector from litellm.llms.bedrock.chat.converse_to_openai import as_json_array, as_json_object +from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + bedrock_invoke_guardrail_headers, +) +from litellm.llms.bedrock.request_metadata import merge_bedrock_invoke_headers from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage from litellm.types.utils import ModelResponseStream @@ -42,6 +46,37 @@ def _is_converse_url(url: str) -> bool: return httpx.URL(url).path.rstrip("/").rsplit("/", 1)[-1] in _CONVERSE_ACTIONS +_INVOKE_ACTIONS: Final = frozenset({"invoke", "invoke-with-response-stream"}) +_BEDROCK_INVOKE_GUARDRAIL_HEADERS: Final = frozenset( + {"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"} +) + + +def _converse_guardrail_config(action: str, guardrail_config: object) -> object: + fields: Final = as_json_object(guardrail_config) + if action == "converse-stream" or fields is None: + return guardrail_config + return {key: value for key, value in fields.items() if key != "streamProcessingMode"} + + +def _with_deployment_guardrail( + action: str, + headers: Mapping[str, str], + body: Mapping[str, object], + guardrail_config: object, +) -> tuple[Mapping[str, str], Mapping[str, object]]: + if guardrail_config is None: + return headers, body + if action in _CONVERSE_ACTIONS: + return headers, {**body, "guardrailConfig": _converse_guardrail_config(action, guardrail_config)} + if action in _INVOKE_ACTIONS: + guardrail_headers: Final = tuple(bedrock_invoke_guardrail_headers(guardrail_config).items()) + return merge_bedrock_invoke_headers( + dict(headers), (), guardrail_headers, _BEDROCK_INVOKE_GUARDRAIL_HEADERS + ), body + return headers, body + + def _is_cache_point(entry: object) -> bool: fields: Final = as_json_object(entry) return fields is not None and "cachePoint" in fields @@ -363,13 +398,20 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD api_base: str, model: str | None = None, ) -> tuple[dict, bytes | None]: + action: Final = httpx.URL(api_base).path.rstrip("/").rsplit("/", 1)[-1] body: Final[Mapping[str, object]] = request_data or {} optional_params: Final = litellm_params.copy() + guarded_headers, guarded_body = _with_deployment_guardrail( + action, + headers, + _converse_body_to_send(body) if action in _CONVERSE_ACTIONS else body, + litellm_params.get("guardrailConfig"), + ) return self._sign_request( service_name="bedrock", - headers=headers, + headers=dict(guarded_headers), optional_params=optional_params, - request_data=dict(_converse_body_to_send(body) if _is_converse_url(api_base) else body), + request_data=dict(guarded_body), api_base=api_base, model=model, api_key=self.get_bedrock_bearer_token(optional_params), diff --git a/litellm/llms/bedrock/request_metadata.py b/litellm/llms/bedrock/request_metadata.py index 1f4e5886508..685b52e3b1c 100644 --- a/litellm/llms/bedrock/request_metadata.py +++ b/litellm/llms/bedrock/request_metadata.py @@ -179,12 +179,11 @@ def merge_bedrock_invoke_headers( """ Merge the ``X-Amzn-*`` headers the Invoke paths derive from params. - ``caller_owned`` (the guardrail headers) defers to a header the caller already set, which is - the long-standing behaviour for those. ``proxy_owned_names`` are dropped from the caller's - headers unconditionally and re-supplied only from ``proxy_owned``, because those names carry - proxy-authenticated identity into an AWS billing record that the caller must not be able to - write. Names are compared case-insensitively so a caller cannot leave a second spelling in - the dict and let the transport pick the winner. + ``caller_owned`` defers to a header the caller already set. ``proxy_owned_names`` are dropped + from the caller's headers unconditionally and re-supplied only from ``proxy_owned``, for names + the caller must not be able to write, such as proxy-authenticated identity headed for an AWS + billing record or a deployment's guardrail. Names are compared case-insensitively so a caller + cannot leave a second spelling in the dict and let the transport pick the winner. """ if not caller_owned and not proxy_owned and not proxy_owned_names: return headers diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 744f9e1aa4c..4e94b1bb910 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -437,6 +437,7 @@ def llm_passthrough_route( params: QueryParamTypes | None = None, cookies: CookieTypes | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, + guardrailConfig: object = None, **kwargs, ) -> ( httpx.Response @@ -547,7 +548,7 @@ def llm_passthrough_route( _request_data: dict | None = data if isinstance(data, dict) else (json if isinstance(json, dict) else None) headers, signed_json_body = provider_config.sign_request( headers=headers, - litellm_params=litellm_params_dict, + litellm_params={**litellm_params_dict, "guardrailConfig": guardrailConfig}, request_data=_request_data, api_base=str(updated_url), model=model, diff --git a/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py b/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py new file mode 100644 index 00000000000..f995a7ac83f --- /dev/null +++ b/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py @@ -0,0 +1,531 @@ +import base64 +import json +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import boto3 +import pytest +import yaml +from botocore.config import Config +from botocore.exceptions import ClientError +from integration._support.budget import case_timeout_seconds +from integration._support.client import Gateway +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 180) + +_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_BLOCKED_WORD: Final = "zzzforbidden" +_BLOCKED_TEXT: Final = "Blocked by guardrail" +_VERSIONLESS_GUARDRAIL_MESSAGE: Final = "guardrailIdentifier and guardrailVersion must be provided together." +_ACTIONS: Final = ("converse", "converse-stream", "invoke", "invoke-with-response-stream") +_DEPLOYMENT_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "deployment-guardrail", + "guardrailVersion": "3", + "trace": "enabled", + "streamProcessingMode": "sync", +} +_CALLER_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "caller-guardrail", + "guardrailVersion": "9", + "trace": "disabled", +} +_MANTLE_KEY: Final = "scripted-mantle-key" +_CALLER_INVOKE_HEADERS: Final[Mapping[str, str]] = { + "x-pass-x-amzn-bedrock-guardrailidentifier": "caller-guardrail", + "x-pass-x-amzn-bedrock-guardrailversion": "9", + "x-pass-x-amzn-bedrock-trace": "DISABLED", +} + + +@dataclass(frozen=True, slots=True) +class Applied: + """The guardrail a request asked Bedrock for, read from where Bedrock reads it for that action.""" + + identifier: str + version: str | None + trace: str | None + stream_processing_mode: str | None + + +@dataclass(frozen=True, slots=True) +class Outcome: + """What the AWS SDK caller reads back.""" + + text: str + intervened: bool + traced_guardrail: str | None + + +def _action(request: Request) -> str: + return request.target.partition("?")[0].rsplit("/", 1)[-1] + + +def _applied(request: Request) -> Applied | None: + if _action(request).startswith("converse"): + config: Final = json.loads(request.body).get("guardrailConfig") + if config is None: + return None + return Applied( + config["guardrailIdentifier"], + config.get("guardrailVersion"), + config.get("trace"), + config.get("streamProcessingMode"), + ) + identifier: Final = request.headers.get("x-amzn-bedrock-guardrailidentifier") + if identifier is None: + return None + return Applied( + identifier, + request.headers.get("x-amzn-bedrock-guardrailversion"), + request.headers.get("x-amzn-bedrock-trace"), + None, + ) + + +def _identifier(request: Request) -> str | None: + applied: Final = _applied(request) + return None if applied is None else applied.identifier + + +def _assessment(applied: Applied) -> Mapping[str, JsonValue]: + return { + applied.identifier: { + "wordPolicy": {"customWords": [{"match": _BLOCKED_WORD, "action": "BLOCKED", "detected": True}]}, + "appliedGuardrailDetails": {"guardrailId": applied.identifier, "guardrailVersion": applied.version}, + } + } + + +def _bedrock_with_guardrails(request: Request) -> Reply: + """Bedrock's runtime as the caller sees it: the guardrail the request names blocks a prompt holding the blocked + word, answering with its message and, when the request enabled tracing, a trace naming that guardrail. A guardrail + named without a version is refused with a 400 ValidationException, as real Bedrock in us-east-1 refused it on + every one of these actions on 2026-10-19.""" + applied: Final = _applied(request) + if applied is not None and applied.version is None: + return Reply( + status=400, + body=json.dumps({"message": _VERSIONLESS_GUARDRAIL_MESSAGE}).encode(), + headers={"x-amzn-errortype": "ValidationException"}, + ) + blocked: Final = applied is not None and _BLOCKED_WORD.encode() in request.body + traced: Final = blocked and applied is not None and (applied.trace or "").lower() == "enabled" + text: Final = _BLOCKED_TEXT if blocked else "ok" + stop_reason: Final = "guardrail_intervened" if blocked else "end_turn" + usage: Final = {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6} + converse_trace: Final = {"trace": {"guardrail": {"inputAssessment": _assessment(applied)}}} if traced else {} + action: Final = _action(request) + if action == "converse": + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": text}]}}, + "stopReason": stop_reason, + "usage": usage, + "metrics": {"latencyMs": 3}, + **converse_trace, + } + ).encode() + ) + if action == "converse-stream": + return Reply( + chunks=( + _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"), + _aws_event_frame("contentBlockDelta", {"delta": {"text": text}, "contentBlockIndex": 0}, "sc", "u"), + _aws_event_frame("contentBlockStop", {"contentBlockIndex": 0}, "sc", "u"), + _aws_event_frame("messageStop", {"stopReason": stop_reason}, "sc", "u"), + _aws_event_frame( + "metadata", {"usage": usage, "metrics": {"latencyMs": 3}, **converse_trace}, "sc", "u" + ), + ), + content_type=_EVENT_STREAM, + ) + message: Final = { + "id": "msg_bedrock", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": text}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, + **({"amazon-bedrock-guardrailAction": "INTERVENED"} if blocked else {}), + **({"amazon-bedrock-trace": {"guardrail": {"input": _assessment(applied)}}} if traced else {}), + } + if action == "invoke": + return Reply(body=json.dumps(message).encode()) + chunk: Final = {"bytes": base64.b64encode(json.dumps(message).encode()).decode()} + return Reply(chunks=(_aws_event_frame("chunk", chunk, "sc", "u"),), content_type=_EVENT_STREAM) + + +def _bedrock_outage(request: Request) -> Reply: + return Reply( + status=500, + body=json.dumps({"message": "scripted outage"}).encode(), + headers={"x-amzn-errortype": "InternalServerException"}, + ) + + +def _traced_guardrail(assessment: object) -> str | None: + return None if not isinstance(assessment, dict) else next(iter(assessment)) + + +def _send(runtime: object, action: str, model: str, prompt: str, caller_guardrail: bool) -> Outcome: + messages: Final = [{"role": "user", "content": [{"text": prompt}]}] + invoke_body: Final = json.dumps( + {"anthropic_version": "bedrock-2023-05-31", "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]} + ) + invoke_guardrail: Final = ( + {"guardrailIdentifier": "caller-guardrail", "guardrailVersion": "9", "trace": "DISABLED"} + if caller_guardrail + else {} + ) + if action == "converse": + response: Final = getattr(runtime, "converse")( + modelId=model, + messages=messages, + **({"guardrailConfig": dict(_CALLER_GUARDRAIL)} if caller_guardrail else {}), + ) + return Outcome( + text=response["output"]["message"]["content"][0]["text"], + intervened=response["stopReason"] == "guardrail_intervened", + traced_guardrail=_traced_guardrail(response.get("trace", {}).get("guardrail", {}).get("inputAssessment")), + ) + if action == "converse-stream": + events: Final = list( + getattr(runtime, "converse_stream")( + modelId=model, + messages=messages, + **( + {"guardrailConfig": {**_CALLER_GUARDRAIL, "streamProcessingMode": "async"}} + if caller_guardrail + else {} + ), + )["stream"] + ) + (stop,) = tuple(event["messageStop"]["stopReason"] for event in events if "messageStop" in event) + (metadata,) = tuple(event["metadata"] for event in events if "metadata" in event) + return Outcome( + text="".join( + event["contentBlockDelta"]["delta"]["text"] for event in events if "contentBlockDelta" in event + ), + intervened=stop == "guardrail_intervened", + traced_guardrail=_traced_guardrail(metadata.get("trace", {}).get("guardrail", {}).get("inputAssessment")), + ) + if action == "invoke": + message: Final = json.loads( + getattr(runtime, "invoke_model")(modelId=model, body=invoke_body, **invoke_guardrail)["body"].read() + ) + else: + (message,) = tuple( + json.loads(event["chunk"]["bytes"]) + for event in getattr(runtime, "invoke_model_with_response_stream")( + modelId=model, body=invoke_body, **invoke_guardrail + )["body"] + ) + return Outcome( + text=message["content"][0]["text"], + intervened=message.get("amazon-bedrock-guardrailAction") == "INTERVENED", + traced_guardrail=_traced_guardrail(message.get("amazon-bedrock-trace", {}).get("guardrail", {}).get("input")), + ) + + +def _signed_header_names(request: Request) -> frozenset[str]: + return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")) + + +def _deployment(wire: Wire, guardrail_config: Mapping[str, str] | None = None) -> Mapping[str, JsonValue]: + return { + "model": f"bedrock/{_MODEL_ID}", + "api_base": wire.url, + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + "use_in_pass_through": True, + **({} if guardrail_config is None else {"guardrailConfig": dict(guardrail_config)}), + } + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + wire: Wire + guarded_model: str + unguarded_model: str + balanced_model: str + balanced_wires: Mapping[str, Wire] + failing_model: str + failing_wire: Wire + fallback_wire: Wire + mantle_model: str + versionless_model: str + key: str + + def runtime( # noqa: ANN201 # boto3 clients are untyped + self, monkeypatch: pytest.MonkeyPatch, extra_headers: Mapping[str, str] = MappingProxyType({}) + ): + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", self.key) + client: Final = boto3.client( + "bedrock-runtime", + endpoint_url=f"{str(self.proxy.client.base_url).rstrip('/')}/bedrock", + region_name="us-east-1", + config=Config(read_timeout=case_timeout_seconds(), retries={"total_max_attempts": 1}), + ) + client.meta.events.register( + "before-send.bedrock-runtime.*", + lambda request, **_: request.headers.update(extra_headers), + ) + return client + + def received(self, wire: Wire, marker: str) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if marker.encode() in request.body) + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + suffix: Final = uuid.uuid4().hex[:12] + guarded_model: Final = f"bedrock-guarded-{suffix}" + unguarded_model: Final = f"bedrock-unguarded-{suffix}" + balanced_model: Final = f"bedrock-balanced-{suffix}" + failing_model: Final = f"bedrock-failing-{suffix}" + fallback_model: Final = f"bedrock-fallback-{suffix}" + mantle_model: Final = f"bedrock-mantle-{suffix}" + versionless_model: Final = f"bedrock-versionless-{suffix}" + directory: Final = tmp_path_factory.mktemp("bedrock-guardrail-config") + base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + with ( + wire_server(_bedrock_with_guardrails) as wire, + wire_server(_bedrock_with_guardrails) as wire_a, + wire_server(_bedrock_with_guardrails) as wire_b, + wire_server(_bedrock_outage) as failing_wire, + wire_server(_bedrock_with_guardrails) as fallback_wire, + ): + config: Final = { + **base_config, + "model_list": [ + *base_config.get("model_list", []), + { + "model_name": guarded_model, + "litellm_params": _deployment(wire, _DEPLOYMENT_GUARDRAIL), + }, + {"model_name": unguarded_model, "litellm_params": _deployment(wire)}, + { + "model_name": balanced_model, + "litellm_params": _deployment( + wire_a, {"guardrailIdentifier": "guardrail-a", "guardrailVersion": "1"} + ), + }, + { + "model_name": balanced_model, + "litellm_params": _deployment( + wire_b, {"guardrailIdentifier": "guardrail-b", "guardrailVersion": "2"} + ), + }, + { + "model_name": failing_model, + "litellm_params": _deployment( + failing_wire, + {"guardrailIdentifier": "failing-guardrail", "guardrailVersion": "5", "trace": "enabled"}, + ), + }, + { + "model_name": fallback_model, + "litellm_params": _deployment( + fallback_wire, + {"guardrailIdentifier": "fallback-guardrail", "guardrailVersion": "6", "trace": "enabled"}, + ), + }, + { + "model_name": versionless_model, + "litellm_params": _deployment( + wire, {"guardrailIdentifier": "versionless-guardrail", "trace": "enabled"} + ), + }, + { + "model_name": mantle_model, + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-sonnet-5", + "api_base": wire.url, + "api_key": _MANTLE_KEY, + "aws_region_name": "us-east-1", + "use_in_pass_through": True, + "guardrailConfig": { + "guardrailIdentifier": "mantle-guardrail", + "guardrailVersion": "4", + "trace": "enabled", + }, + }, + }, + ], + "router_settings": { + **base_config.get("router_settings", {}), + "fallbacks": [{failing_model: [fallback_model]}], + }, + } + config_path: Final = directory / "bedrock-guardrail-config.yaml" + config_path.write_text(yaml.safe_dump(config)) + with ( + owned_proxy(session_gateway, directory, {}, config=config_path, workers=2) as proxy, + proxy.scenario() as scenario, + ): + yield Rig( + proxy=proxy, + wire=wire, + guarded_model=guarded_model, + unguarded_model=unguarded_model, + balanced_model=balanced_model, + balanced_wires={"guardrail-a": wire_a, "guardrail-b": wire_b}, + failing_model=failing_model, + failing_wire=failing_wire, + fallback_wire=fallback_wire, + mantle_model=mantle_model, + versionless_model=versionless_model, + key=scenario.key(key_alias=f"guardrail-key-{suffix}"), + ) + + +def _expected_on_the_wire(action: str, request: Request) -> Applied: + if action == "converse": + return Applied("deployment-guardrail", "3", "enabled", None) + if action == "converse-stream": + return Applied("deployment-guardrail", "3", "enabled", "sync") + assert {"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"} <= ( + _signed_header_names(request) + ), request.headers["authorization"] + return Applied("deployment-guardrail", "3", "ENABLED", None) + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_passthrough_on_a_guardrailed_deployment_blocks_with_the_deployment_guardrail( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + blocked: Final = _send(runtime, action, rig.guarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False) + allowed: Final = _send(runtime, action, rig.guarded_model, f"{marker} say hello", caller_guardrail=False) + + assert blocked == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="deployment-guardrail"), blocked + assert allowed == Outcome("ok", intervened=False, traced_guardrail=None), allowed + received: Final = rig.received(rig.wire, marker) + assert [_action(request) for request in received] == [action, action], received + assert [_applied(request) for request in received] == [ + _expected_on_the_wire(action, request) for request in received + ], [(request.headers, request.body) for request in received] + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_caller_guardrail_is_replaced_by_the_deployment_guardrail( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch, _CALLER_INVOKE_HEADERS) + + outcome: Final = _send(runtime, action, rig.guarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=True) + + assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="deployment-guardrail"), outcome + (request,) = rig.received(rig.wire, marker) + assert _applied(request) == _expected_on_the_wire(action, request), (request.headers, request.body) + assert "caller-guardrail" not in json.dumps({**request.headers, "body": request.body.decode()}), request + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_mantle_deployment_passthrough_sends_its_guardrail_with_bearer_auth( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcome: Final = _send(runtime, action, rig.mantle_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False) + + assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="mantle-guardrail"), outcome + (request,) = rig.received(rig.wire, marker) + assert request.headers["authorization"] == f"Bearer {_MANTLE_KEY}", request.headers + assert _applied(request) == Applied( + "mantle-guardrail", "4", "ENABLED" if action.startswith("invoke") else "enabled", None + ), (request.headers, request.body) + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_deployment_guardrail_without_a_version_returns_the_bedrock_rejection( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + with pytest.raises(ClientError) as raised: + _send(runtime, action, rig.versionless_model, f"{marker} say hello", caller_guardrail=False) + + error: Final = raised.value.response.get("Error", {}) + assert (error.get("Code"), error.get("Message")) == ("ValidationException", _VERSIONLESS_GUARDRAIL_MESSAGE), error + received: Final = rig.received(rig.wire, marker) + assert {_applied(request) for request in received} == { + Applied("versionless-guardrail", None, "ENABLED" if action.startswith("invoke") else "enabled", None) + }, [(request.headers, request.body) for request in received] + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_deployment_without_a_guardrail_leaves_the_caller_guardrail_as_sent( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + if action.startswith("invoke"): + pytest.skip("BUG: the router-model invoke pass-through drops the caller's X-Amzn-Bedrock-Guardrail* headers") + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcome: Final = _send(runtime, action, rig.unguarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=True) + + assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail=None), outcome + (request,) = rig.received(rig.wire, marker) + assert _applied(request) == Applied( + "caller-guardrail", "9", "disabled", "async" if action == "converse-stream" else None + ), request.body + + +def test_each_load_balanced_request_carries_the_guardrail_of_the_deployment_it_reached( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcomes: Final = tuple( + _send(runtime, action, rig.balanced_model, f"{marker}-{attempt} say {_BLOCKED_WORD}", caller_guardrail=False) + for attempt, action in enumerate(_ACTIONS * 6) + ) + + assert all(outcome.intervened for outcome in outcomes), outcomes + reached: Final = {name: rig.received(wire, marker) for name, wire in rig.balanced_wires.items()} + assert sum(len(requests) for requests in reached.values()) == len(outcomes), reached + assert all(reached.values()), {name: len(requests) for name, requests in reached.items()} + assert {name: {_identifier(request) for request in requests} for name, requests in reached.items()} == { + "guardrail-a": {"guardrail-a"}, + "guardrail-b": {"guardrail-b"}, + }, reached + + +def test_a_fallback_sends_the_guardrail_of_the_deployment_it_fell_back_to( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcomes: Final = tuple( + _send(runtime, action, rig.failing_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False) + for action in _ACTIONS + ) + + assert outcomes == (Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="fallback-guardrail"),) * 4, outcomes + failed: Final = rig.received(rig.failing_wire, marker) + served: Final = rig.received(rig.fallback_wire, marker) + assert [_action(request) for request in served] == list(_ACTIONS), served + assert {_identifier(request) for request in failed} == {"failing-guardrail"}, failed + assert {_identifier(request) for request in served} == {"fallback-guardrail"}, served diff --git a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index 84d3160e609..4b7e7aad5e4 100644 --- a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -5,8 +5,9 @@ import struct import sys import types from binascii import crc32 -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime +from types import MappingProxyType from typing import Final from unittest.mock import patch @@ -17,6 +18,7 @@ import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.passthrough.transformation import ( BedrockPassthroughConfig, is_bedrock_streaming_endpoint, @@ -898,3 +900,194 @@ def test_stream_collector_logs_nothing_for_an_unrecognized_endpoint(): collector.add(_event_frame("messageStart", {"role": "assistant"})) assert collector.build_logged_response(_converse_stream_logging_obj()) is None + + +CONVERSE_BODY = {"messages": [{"role": "user", "content": [{"text": "hi"}]}]} +INVOKE_BODY = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], +} +DEPLOYMENT_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "deployment-guardrail", + "guardrailVersion": "3", + "trace": "enabled", + "streamProcessingMode": "sync", +} +DEPLOYMENT_INVOKE_GUARDRAIL_HEADERS: Final[Mapping[str, str]] = { + "x-amzn-bedrock-guardrailidentifier": "deployment-guardrail", + "x-amzn-bedrock-guardrailversion": "3", + "x-amzn-bedrock-trace": "ENABLED", +} +CALLER_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "caller-guardrail", + "guardrailVersion": "9", + "trace": "disabled", +} +CALLER_INVOKE_GUARDRAIL_HEADERS: Final[Mapping[str, str]] = { + "x-pass-X-Amzn-Bedrock-GuardrailIdentifier": "caller-guardrail", + "x-pass-X-Amzn-Bedrock-GuardrailVersion": "9", + "x-pass-X-Amzn-Bedrock-Trace": "DISABLED", +} + + +async def _sent_passthrough_request( + action: str, + body: dict[str, object], + deployment_params: Mapping[str, object] = MappingProxyType({}), + **kwargs: object, +) -> httpx.Request: + requests: Final[list[httpx.Request]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, json={"ok": True}, request=request) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "bedrock-pt", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_access_key_id": "AKIAROUTERKEY", + "aws_secret_access_key": "router-secret", + "aws_region_name": "us-west-2", + **deployment_params, + }, + } + ] + ) + await router.allm_passthrough_route( + model="bedrock-pt", + method="POST", + endpoint=f"/model/bedrock-pt/{action}", + data=body, + client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)), + **kwargs, + ) + return requests[0] + + +def _signed_header_names(request: httpx.Request) -> frozenset[str]: + return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")) + + +def _invoke_guardrail_headers(request: httpx.Request) -> dict[str, str]: + return { + name: value + for name, value in request.headers.items() + if name.lower().startswith("x-amzn-bedrock-guardrail") or name.lower() == "x-amzn-bedrock-trace" + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "sent_guardrail"), + [ + ("converse", {"guardrailIdentifier": "deployment-guardrail", "guardrailVersion": "3", "trace": "enabled"}), + ("converse-stream", DEPLOYMENT_GUARDRAIL), + ], +) +async def test_converse_passthrough_sends_the_deployment_guardrail_in_the_body( + action: str, sent_guardrail: Mapping[str, str] +) -> None: + request: Final = await _sent_passthrough_request( + action, CONVERSE_BODY, deployment_params={"guardrailConfig": DEPLOYMENT_GUARDRAIL} + ) + + assert json.loads(request.content) == {**CONVERSE_BODY, "guardrailConfig": sent_guardrail} + assert _invoke_guardrail_headers(request) == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_passthrough_sends_the_deployment_guardrail_as_signed_headers(action: str) -> None: + request: Final = await _sent_passthrough_request( + action, INVOKE_BODY, deployment_params={"guardrailConfig": DEPLOYMENT_GUARDRAIL} + ) + + assert _invoke_guardrail_headers(request) == DEPLOYMENT_INVOKE_GUARDRAIL_HEADERS + assert set(DEPLOYMENT_INVOKE_GUARDRAIL_HEADERS) <= _signed_header_names(request) + assert json.loads(request.content) == INVOKE_BODY + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "caller_guardrail"), + [("converse", CALLER_GUARDRAIL), ("converse-stream", {**CALLER_GUARDRAIL, "streamProcessingMode": "async"})], +) +async def test_deployment_guardrail_replaces_the_caller_converse_guardrail( + action: str, caller_guardrail: Mapping[str, str] +) -> None: + request: Final = await _sent_passthrough_request( + action, + {**CONVERSE_BODY, "guardrailConfig": caller_guardrail}, + deployment_params={"guardrailConfig": {"guardrailIdentifier": "deployment-guardrail"}}, + ) + + assert json.loads(request.content) == { + **CONVERSE_BODY, + "guardrailConfig": {"guardrailIdentifier": "deployment-guardrail"}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_deployment_guardrail_replaces_every_invoke_guardrail_header_the_caller_sent(action: str) -> None: + request: Final = await _sent_passthrough_request( + action, + INVOKE_BODY, + deployment_params={"guardrailConfig": {"guardrailIdentifier": "deployment-guardrail"}}, + request_headers=CALLER_INVOKE_GUARDRAIL_HEADERS, + ) + + assert _invoke_guardrail_headers(request) == {"x-amzn-bedrock-guardrailidentifier": "deployment-guardrail"} + assert json.loads(request.content) == INVOKE_BODY + + +@pytest.mark.asyncio +async def test_passthrough_without_a_deployment_guardrail_sends_the_caller_guardrail_as_is() -> None: + caller_body: Final = {**CONVERSE_BODY, "guardrailConfig": CALLER_GUARDRAIL} + + converse: Final = await _sent_passthrough_request("converse", caller_body) + invoke: Final = await _sent_passthrough_request("invoke", INVOKE_BODY, request_headers=CALLER_INVOKE_GUARDRAIL_HEADERS) + + assert json.loads(converse.content) == caller_body + assert _invoke_guardrail_headers(invoke) == { + "x-amzn-bedrock-guardrailidentifier": "caller-guardrail", + "x-amzn-bedrock-guardrailversion": "9", + "x-amzn-bedrock-trace": "DISABLED", + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +@pytest.mark.parametrize("guardrail_config", ["deployment-guardrail", {"guardrailVersion": "3"}]) +async def test_invoke_passthrough_refuses_a_deployment_guardrail_it_cannot_send( + action: str, guardrail_config: object +) -> None: + with pytest.raises(BedrockError, match="guardrailConfig") as rejected: + await _sent_passthrough_request(action, INVOKE_BODY, deployment_params={"guardrailConfig": guardrail_config}) + + assert rejected.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_converse_passthrough_forwards_a_non_object_deployment_guardrail_unchanged() -> None: + request: Final = await _sent_passthrough_request( + "converse", CONVERSE_BODY, deployment_params={"guardrailConfig": "deployment-guardrail"} + ) + + assert json.loads(request.content) == {**CONVERSE_BODY, "guardrailConfig": "deployment-guardrail"} + + +@pytest.mark.asyncio +async def test_passthrough_leaves_actions_without_a_guardrail_field_as_sent() -> None: + count_tokens_body: Final = {"input": {"converse": CONVERSE_BODY}} + + request: Final = await _sent_passthrough_request( + "count-tokens", count_tokens_body, deployment_params={"guardrailConfig": DEPLOYMENT_GUARDRAIL} + ) + + assert json.loads(request.content) == count_tokens_body + assert _invoke_guardrail_headers(request) == {} From 4fb2f026a03bb9d676f9fad1c214a36d0244a209 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 7 Nov 2026 16:58:22 -0800 Subject: [PATCH 37/41] refactor(bedrock): drop the unused converse URL helper and cover a trailing-slash converse pass-through --- litellm/llms/bedrock/passthrough/transformation.py | 4 ---- .../passthrough/test_bedrock_passthrough_transformation.py | 5 ++++- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index 49eea6f5daf..453e730e7fc 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -42,10 +42,6 @@ def is_bedrock_streaming_endpoint(endpoint: str) -> bool: _CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) -def _is_converse_url(url: str) -> bool: - return httpx.URL(url).path.rstrip("/").rsplit("/", 1)[-1] in _CONVERSE_ACTIONS - - _INVOKE_ACTIONS: Final = frozenset({"invoke", "invoke-with-response-stream"}) _BEDROCK_INVOKE_GUARDRAIL_HEADERS: Final = frozenset( {"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"} diff --git a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index 4b7e7aad5e4..e418f6778ee 100644 --- a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -986,6 +986,7 @@ def _invoke_guardrail_headers(request: httpx.Request) -> dict[str, str]: [ ("converse", {"guardrailIdentifier": "deployment-guardrail", "guardrailVersion": "3", "trace": "enabled"}), ("converse-stream", DEPLOYMENT_GUARDRAIL), + ("converse/", {"guardrailIdentifier": "deployment-guardrail", "guardrailVersion": "3", "trace": "enabled"}), ], ) async def test_converse_passthrough_sends_the_deployment_guardrail_in_the_body( @@ -1050,7 +1051,9 @@ async def test_passthrough_without_a_deployment_guardrail_sends_the_caller_guard caller_body: Final = {**CONVERSE_BODY, "guardrailConfig": CALLER_GUARDRAIL} converse: Final = await _sent_passthrough_request("converse", caller_body) - invoke: Final = await _sent_passthrough_request("invoke", INVOKE_BODY, request_headers=CALLER_INVOKE_GUARDRAIL_HEADERS) + invoke: Final = await _sent_passthrough_request( + "invoke", INVOKE_BODY, request_headers=CALLER_INVOKE_GUARDRAIL_HEADERS + ) assert json.loads(converse.content) == caller_body assert _invoke_guardrail_headers(invoke) == { From d06f5a5021ed67982c6a19302470b9b790554b93 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 7 Nov 2026 17:07:56 -0800 Subject: [PATCH 38/41] test(e2e): repin the Together reasoning replay backend to MiniMax-M3.1-Flash MiniMax-M3 reached its 2026-11-06 deprecation date in the cost map, which names MiniMax-M3.1-Flash as its successor (cherry picked from commit 90f7ab6fe6837152b21617c8cb7dc73a541c42e9) --- tests/e2e/llm_translation/test_together_ai_e2e.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/e2e/llm_translation/test_together_ai_e2e.py b/tests/e2e/llm_translation/test_together_ai_e2e.py index 2d5d5f2d93a..7775c5c6ed2 100644 --- a/tests/e2e/llm_translation/test_together_ai_e2e.py +++ b/tests/e2e/llm_translation/test_together_ai_e2e.py @@ -7,9 +7,9 @@ backends are pinned because the registry has no flag for what they prove: ``enab and the ``{"reasoning": {"enabled": false}}`` toggle that ``reasoning_effort="none"`` maps to are Qwen hybrid-model contracts, and the structured-output case runs on that hybrid model with reasoning off, since a reasoning-only model can spend the whole token budget -thinking and return no content. MiniMax-M3 is the serverless model whose +thinking and return no content. MiniMax-M3.1-Flash is the serverless model whose template renders a replayed ``reasoning_content`` back into the prompt (Qwen and -DeepSeek silently drop it). MiniMax-M3 honors that replayed field on nearly every call, not +DeepSeek silently drop it). MiniMax-M3.1-Flash honors that replayed field on nearly every call, not every call (one miss in dozens of otherwise identical calls), so the replay case asks up to ``REPLAY_ATTEMPTS`` times and fails only when no answer carries the secret, which a proxy that strips the field guarantees. Requires TOGETHER_API_KEY on the proxy; no @@ -56,7 +56,7 @@ from pydantic import BaseModel pytestmark = pytest.mark.e2e HYBRID_REASONING_BACKEND = "together_ai/Qwen/Qwen3.5-9B" -REASONING_REPLAY_BACKEND = "together_ai/MiniMaxAI/MiniMax-M3" +REASONING_REPLAY_BACKEND = "together_ai/MiniMaxAI/MiniMax-M3.1-Flash" SECRET_PROMPT = "Remember this for later and reply with just OK." SECRET_REASONING = "The user told me their favorite color is chartreuse. I must remember it." From cb738a569765fce2ebdd6c1f8a423f5611e2eff3 Mon Sep 17 00:00:00 2001 From: MFJG <59931828+MFJG-dev@users.noreply.github.com> Date: Mon, 19 Oct 2026 21:44:54 -0400 Subject: [PATCH 39/41] fix(sdk): backport #45961 to lts/1.108.x completion() loaded the MCP gateway handler for any request with tools, and the handler pulls in mcp 2.x only code, so tool calls failed on installs with mcp 1.x (google-adk 1.36.1) or no mcp at all Cherry-picked from 783cf06f31. Only conflict was import context in litellm/responses/main.py (cherry picked from commit 783cf06f31e55e8f7ab84c2d7c708b0d90e68001) --- .../pass_through/messages/handler.py | 53 ++- litellm/main.py | 103 +++--- litellm/responses/main.py | 7 +- .../mcp/litellm_proxy_mcp_handler.py | 44 +-- litellm/responses/mcp/tool_precheck.py | 43 +++ .../base_sdk_tests/check_base_sdk_install.py | 24 ++ tests/base_sdk_tests/check_sdk_http.py | 55 ++- .../integration/mcp/test_mcp_llm_endpoints.py | 55 +++ .../unit/responses/mcp/test_tool_precheck.py | 318 ++++++++++++++++++ .../test_responses_prompt_management.py | 35 +- 10 files changed, 593 insertions(+), 144 deletions(-) create mode 100644 litellm/responses/mcp/tool_precheck.py create mode 100644 tests/unit/responses/mcp/test_tool_precheck.py diff --git a/litellm/llms/anthropic/pass_through/messages/handler.py b/litellm/llms/anthropic/pass_through/messages/handler.py index eeff6302ee8..949d81f3c1b 100644 --- a/litellm/llms/anthropic/pass_through/messages/handler.py +++ b/litellm/llms/anthropic/pass_through/messages/handler.py @@ -27,6 +27,7 @@ from litellm.llms.base_llm.anthropic_messages.transformation import ( from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway from litellm.types.llms.anthropic import AllAnthropicToolsValues from litellm.types.llms.anthropic_messages.anthropic_request import AnthropicMetadata from litellm.types.llms.anthropic_messages.anthropic_response import ( @@ -569,38 +570,34 @@ def anthropic_messages_handler( # downstream path (native passthrough and both bridges) gets real tools rather than a # reference the provider cannot resolve. Popped from kwargs so it never reaches the provider. skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False) - if not skip_mcp_handler and tools: + if not skip_mcp_handler and tools and may_name_litellm_mcp_gateway(tools): # pyright: ignore[reportUnknownArgumentType] # untyped tools list from litellm.llms.anthropic.pass_through.messages.mcp_handler import ( anthropic_messages_with_mcp, ) - from litellm.responses.mcp.litellm_proxy_mcp_handler import ( - LiteLLM_Proxy_MCP_Handler, - ) - if LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(tools=tools): - return anthropic_messages_with_mcp( - max_tokens=max_tokens, - messages=messages, - model=model, - metadata=metadata, - stop_sequences=stop_sequences, - stream=stream, - system=cast( # cast-ok: anthropic accepts list system blocks; handler signature is narrower - "str | None", system - ), - temperature=temperature, - thinking=thinking, - tool_choice=tool_choice, - tools=tools, - top_k=top_k, - top_p=top_p, - container=container, - api_key=api_key, - api_base=api_base, - client=client, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) + return anthropic_messages_with_mcp( + max_tokens=max_tokens, + messages=messages, + model=model, + metadata=metadata, + stop_sequences=stop_sequences, + stream=stream, + system=cast( # cast-ok: anthropic accepts list system blocks; handler signature is narrower + "str | None", system + ), + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + container=container, + api_key=api_key, + api_base=api_base, + client=client, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) anthropic_messages_provider_config: Final = _native_messages_config( model, custom_llm_provider, kwargs.get("model_info") diff --git a/litellm/main.py b/litellm/main.py index 42d5ecc8f20..fa628cd9aff 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -135,6 +135,7 @@ from litellm.llms.vertex_ai.common_utils import ( get_vertex_ai_model_route, ) from litellm.realtime_api.main import realtime_health_check +from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway from litellm.secret_managers.main import get_secret_bool, get_secret_str from litellm.types.completion import ( CompletionDispatchContext, @@ -5366,63 +5367,55 @@ def completion( control_options: Final = _resolve_control_options(kwargs, model) skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False) - if not skip_mcp_handler and tools: + if not skip_mcp_handler and tools and may_name_litellm_mcp_gateway(tools): # pyright: ignore[reportUnknownArgumentType] # untyped tools list from litellm.responses.mcp.chat_completions_handler import acompletion_with_mcp - from litellm.responses.mcp.litellm_proxy_mcp_handler import ( - LiteLLM_Proxy_MCP_Handler, - ) - from litellm.types.llms.openai import ToolParam - # Check if MCP tools are present (following responses pattern) - # Cast tools to Optional[Iterable[ToolParam]] for type checking - tools_for_mcp: Final = cast(Iterable[ToolParam] | None, tools) - if LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(tools=tools_for_mcp): - return acompletion_with_mcp( # pyright: ignore[reportReturnType] # MCP path returns a coroutine that acompletion() awaits; completion()'s sync return type omits it - model=model, - messages=messages, - functions=functions, - function_call=function_call, - timeout=timeout, - temperature=temperature, - top_p=top_p, - n=n, - stream=stream, - stream_options=stream_options, - stop=stop, - max_tokens=max_tokens, - max_completion_tokens=max_completion_tokens, - modalities=modalities, - prediction=prediction, - audio=audio, - presence_penalty=presence_penalty, - frequency_penalty=frequency_penalty, - logit_bias=logit_bias, - user=user, - response_format=response_format, - seed=seed, - tools=tools, - tool_choice=tool_choice, - parallel_tool_calls=parallel_tool_calls, - logprobs=logprobs, - top_logprobs=top_logprobs, - deployment_id=deployment_id, - reasoning_effort=reasoning_effort, - verbosity=verbosity, - safety_identifier=safety_identifier, - service_tier=service_tier, - store=store, - prompt_cache_key=prompt_cache_key, - base_url=base_url, - api_version=api_version, - api_key=api_key, - model_list=model_list, - extra_headers=extra_headers, - thinking=thinking, - web_search_options=web_search_options, - shared_session=shared_session, - enable_json_schema_validation=enable_json_schema_validation, - **kwargs, - ) + return acompletion_with_mcp( # pyright: ignore[reportReturnType] # MCP path returns a coroutine that acompletion() awaits; completion()'s sync return type omits it + model=model, + messages=messages, + functions=functions, + function_call=function_call, + timeout=timeout, + temperature=temperature, + top_p=top_p, + n=n, + stream=stream, + stream_options=stream_options, + stop=stop, + max_tokens=max_tokens, + max_completion_tokens=max_completion_tokens, + modalities=modalities, + prediction=prediction, + audio=audio, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + logit_bias=logit_bias, + user=user, + response_format=response_format, + seed=seed, + tools=tools, + tool_choice=tool_choice, + parallel_tool_calls=parallel_tool_calls, + logprobs=logprobs, + top_logprobs=top_logprobs, + deployment_id=deployment_id, + reasoning_effort=reasoning_effort, + verbosity=verbosity, + safety_identifier=safety_identifier, + service_tier=service_tier, + store=store, + prompt_cache_key=prompt_cache_key, + base_url=base_url, + api_version=api_version, + api_key=api_key, + model_list=model_list, + extra_headers=extra_headers, + thinking=thinking, + web_search_options=web_search_options, + shared_session=shared_session, + enable_json_schema_validation=enable_json_schema_validation, + **kwargs, + ) api_base = kwargs.get("api_base", None) mock_response: Final[MOCK_RESPONSE_TYPE | None] = kwargs.get("mock_response", None) mock_tool_calls: Final = kwargs.get("mock_tool_calls", None) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 1948597aaf1..a858b91c4c8 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -41,6 +41,7 @@ from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) from litellm.responses.mcp.request_context import MCPRequestContext +from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PromptObject, @@ -1112,11 +1113,7 @@ def _responses_try_dispatch_mcp_gateway( | None ): """Return a response when MCP gateway handles the call; otherwise None.""" - from litellm.responses.mcp.litellm_proxy_mcp_handler import ( - LiteLLM_Proxy_MCP_Handler, - ) - - if skip_mcp_handler or not LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(tools=tools): + if skip_mcp_handler or not may_name_litellm_mcp_gateway(tools): return None mcp_call_kwargs: Final = { "model": model, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 1dfa4b2abcb..e8a9213d4cf 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,4 +1,3 @@ -import re import traceback from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from datetime import datetime @@ -20,6 +19,12 @@ from litellm.proxy._experimental.mcp_server.utils import ( strip_known_server_prefix, ) from litellm.responses.main import aresponses # noqa: TID251 # inner call must skip the MCP gateway that invoked it +from litellm.responses.mcp.tool_precheck import ( + LITELLM_PROXY_MCP_SERVER_URL_PREFIX, + may_name_litellm_mcp_gateway, + names_gateway_explicitly, + proxy_path_mcp_name, +) from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( ResponseInputParam, @@ -65,29 +70,6 @@ class _AresponsesOptions(ExtraItemsTypedDict, total=False, extra_items=object): pass -LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy" -LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" - -_PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$") - - -def _mcp_server_url(tool: ToolParam) -> str | None: - if not isinstance(tool, dict) or tool.get("type") != "mcp": - return None - server_url: Final = tool.get("server_url") - return server_url if isinstance(server_url, str) else None - - -def _names_gateway_explicitly(tool: ToolParam) -> bool: - return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL) - - -def _proxy_path_mcp_name(tool: ToolParam) -> str | None: - server_url: Final = _mcp_server_url(tool) - match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url) - return None if match is None else match.group(1) - - def _registered_mcp_servers() -> Collection[MCPServer]: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, @@ -133,7 +115,7 @@ async def _gateway_served_names( async def _served_mcp_path_names( tools: Collection[ToolParam], served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] ) -> frozenset[str]: - names: Final = frozenset(name for name in map(_proxy_path_mcp_name, tools) if name is not None) + names: Final = frozenset(name for name in map(proxy_path_mcp_name, tools) if name is not None) return await served_names(names) if names else frozenset[str]() @@ -162,17 +144,15 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool: - """True when a tool may name this gateway: server_url "litellm_proxy..." or an http(s) URL ending in - /mcp/. `_split_mcp_tools` then settles which of the latter the gateway actually serves.""" - return any(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) is not None for tool in tools or ()) + return may_name_litellm_mcp_gateway(tools) _should_use_litellm_mcp_gateway = should_use_litellm_mcp_gateway @staticmethod def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools: items: Final = tuple(tools or ()) - gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)] - other_tools: Final[list[ToolParam]] = [tool for tool in items if not _names_gateway_explicitly(tool)] + gateway_tools: Final[list[ToolParam]] = [tool for tool in items if names_gateway_explicitly(tool)] + other_tools: Final[list[ToolParam]] = [tool for tool in items if not names_gateway_explicitly(tool)] return gateway_tools, other_tools @staticmethod @@ -185,7 +165,7 @@ class LiteLLM_Proxy_MCP_Handler: return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools( [ {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"} - if (name := _proxy_path_mcp_name(tool)) in served + if (name := proxy_path_mcp_name(tool)) in served else tool for tool in items ] @@ -201,7 +181,7 @@ class LiteLLM_Proxy_MCP_Handler: ) -> tuple[bool, ...]: items: Final = tuple(tools or ()) served: Final = await _served_mcp_path_names(items, served_names) - return tuple(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) in served for tool in items) + return tuple(names_gateway_explicitly(tool) or proxy_path_mcp_name(tool) in served for tool in items) @staticmethod async def _apply_toolset_permissions( diff --git a/litellm/responses/mcp/tool_precheck.py b/litellm/responses/mcp/tool_precheck.py new file mode 100644 index 00000000000..b76f4066e6d --- /dev/null +++ b/litellm/responses/mcp/tool_precheck.py @@ -0,0 +1,43 @@ +import importlib +import re +from collections.abc import Iterable +from typing import Final + +LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy" +LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" + +_PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$") +_GATEWAY_REGISTRY_MODULE: Final = "litellm.proxy._experimental.mcp_server.mcp_server_manager" + + +def _mcp_server_url(tool: object) -> str | None: + match tool: + case {"type": "mcp", "server_url": str() as server_url} if isinstance(tool, dict): + return server_url + case _: + return None + + +def names_gateway_explicitly(tool: object) -> bool: + return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL) + + +def proxy_path_mcp_name(tool: object) -> str | None: + server_url: Final = _mcp_server_url(tool) + match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url) + return None if match is None else match.group(1) + + +def _gateway_registry_loads() -> bool: + try: + importlib.import_module(_GATEWAY_REGISTRY_MODULE) + except ImportError: + return False + return True + + +def may_name_litellm_mcp_gateway(tools: Iterable[object] | None) -> bool: + items: Final = tuple(tools or ()) + if any(map(names_gateway_explicitly, items)): + return True + return any(proxy_path_mcp_name(tool) is not None for tool in items) and _gateway_registry_loads() diff --git a/tests/base_sdk_tests/check_base_sdk_install.py b/tests/base_sdk_tests/check_base_sdk_install.py index 3abea2a9707..4abd71e39a6 100644 --- a/tests/base_sdk_tests/check_base_sdk_install.py +++ b/tests/base_sdk_tests/check_base_sdk_install.py @@ -102,6 +102,29 @@ def check_completion() -> str: return "mock completion round-trips" +def check_completion_with_function_tool() -> str: + import litellm + + response: Final = litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "ping"}], + tools=[ + { + "type": "function", + "function": { + "name": "noop", + "description": "No operation", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + mock_response="pong", + ) + content: Final = response.choices[0].message.content + _require(content == "pong", f"mock completion with a function tool returned {content!r}") + return "mock completion with a function tool round-trips" + + def check_mcp_install_guidance() -> str: try: import litellm.experimental_mcp_client @@ -209,6 +232,7 @@ CHECKS: tuple[tuple[str, Callable[[], str]], ...] = ( ("import litellm", check_import), ("optional MCP installation guidance", check_mcp_install_guidance), ("chat completion", check_completion), + ("chat completion with a function tool", check_completion_with_function_tool), ("embedding", check_embedding), ("bundled model metadata", check_bundled_model_metadata), ("token counter", check_token_counter), diff --git a/tests/base_sdk_tests/check_sdk_http.py b/tests/base_sdk_tests/check_sdk_http.py index c93306cdced..f7a1e30ac99 100644 --- a/tests/base_sdk_tests/check_sdk_http.py +++ b/tests/base_sdk_tests/check_sdk_http.py @@ -89,7 +89,6 @@ def payload(response: object) -> Mapping[str, JsonValue]: return JSON_OBJECT.validate_json(response.model_dump_json(exclude_none=True)) -@skipIf(importlib.util.find_spec("mcp") is None, "BUG: LIT-9333 completion(tools=...) imports optional mcp") def check_tools(base: str, asynchronous: bool = False) -> None: tools: Final = [ { @@ -122,7 +121,6 @@ def check_tools(base: str, asynchronous: bool = False) -> None: assert json.loads(body) == {"model": "gpt-6", "messages": MESSAGES, "tools": tools} -@skipIf(importlib.util.find_spec("mcp") is None, "BUG: LIT-9333 responses() imports optional mcp") def check_responses(base: str, asynchronous: bool = False) -> None: output: Final[Sequence[JsonValue]] = [ { @@ -144,14 +142,62 @@ def check_responses(base: str, asynchronous: bool = False) -> None: "usage": {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5}, } ) - arguments: Final = dict(model="openai/gpt-6", input="ping", api_key="test-key", api_base=base + "/v1") + tools: Final = [ + {"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com/sse"}, + {"type": "mcp", "server_label": "repos", "server_url": "https://mcp.example.com/mcp/repos"}, + ] + arguments: Final = dict(model="openai/gpt-6", input="ping", tools=tools, api_key="test-key", api_base=base + "/v1") response: Final = asyncio.run(litellm.aresponses(**arguments)) if asynchronous else litellm.responses(**arguments) assert payload(response)["output"] == output assert payload(response)["usage"] == {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5} path, headers, body = REQUESTS.get(timeout=10) assert path == "/v1/responses" assert headers["authorization"] == "Bearer test-key" - assert json.loads(body) == {"model": "gpt-6", "input": "ping"} + assert json.loads(body) == {"model": "gpt-6", "input": "ping", "tools": tools} + + +def check_messages_tools(base: str) -> None: + tools: Final = [ + { + "name": "lookup", + "description": "Look up the forecast for a city", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } + ] + content: Final = [{"type": "tool_use", "id": "toolu-http-check", "name": "lookup", "input": {"city": "Paris"}}] + enqueue( + { + "id": "msg-tools-http-check", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": content, + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 2}, + } + ) + + async def send() -> object: + create: Final = cast( # cast-ok: SDK facade exposes untyped kwargs + Callable[..., Awaitable[object]], litellm.anthropic.messages.acreate + ) + return await create( + model="anthropic/claude-sonnet-5-5", + messages=MESSAGES, + tools=tools, + max_tokens=16, + api_key="test-key", + api_base=base, + ) + + response: Final = JSON_OBJECT.validate_python(asyncio.run(send())) + assert response["content"] == content, response + path, headers, body = REQUESTS.get(timeout=10) + assert path == "/v1/messages" + assert headers["x-api-key"] == "test-key" + sent: Final = JSON_OBJECT.validate_json(body) + assert (sent["model"], sent["messages"], sent["tools"]) == ("claude-sonnet-5-5", MESSAGES, tools), sent @skipIf(importlib.util.find_spec("fastapi") is None, "BUG: LIT-9333 cache hits import proxy-only fastapi") @@ -392,6 +438,7 @@ def main() -> None: partial(check_tools, asynchronous=True), check_responses, partial(check_responses, asynchronous=True), + check_messages_tools, check_cache_hit, check_messages_stream_and_callback, ): diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index c3bccec1f8e..c1517a48c1c 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -468,6 +468,61 @@ def test_server_scoped_gateway_url_exposes_only_that_servers_tools(gateway: Gate assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"] +@pytest.mark.parametrize("surface", ("chat", "responses", "messages")) +def test_proxy_path_url_of_a_served_server_is_listed_and_executed_by_the_gateway( + gateway: Gateway, surface: Surface +) -> None: + with _rig(gateway, surface) as rig: + server_url: Final = f"{str(gateway.client.base_url).rstrip('/')}/mcp/{rig.alias}" + response: Final = rig.send(_granted_key(rig), [{**AUTO, "server_url": server_url}]) + assert response.status_code == 200, response.text + requests: Final = rig.upstream_tools() + assert len(requests) == 2 and all(rig.tool in names for names in requests), requests + peer_calls: Final = tuple(JSON_OBJECT.validate_python(call["body"]) for call in _peer_add_calls(rig.peer)) + assert [object_value(call["params"])["name"] for call in peer_calls] == ["add"], peer_calls + assert rig.final_text(JSON_OBJECT.validate_json(response.content)) == ANSWER, response.text + + +LOOKUP_PARAMETERS: Final = {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]} +CALLER_FUNCTION_TOOLS: Final[Mapping[Surface, Mapping[str, object]]] = { + "chat": {"type": "function", "function": {"name": "lookup", "parameters": LOOKUP_PARAMETERS}}, + "responses": {"type": "function", "name": "lookup", "parameters": LOOKUP_PARAMETERS}, + "messages": {"name": "lookup", "input_schema": LOOKUP_PARAMETERS}, +} + + +@pytest.mark.parametrize("surface", ("chat", "responses", "messages")) +def test_tools_that_do_not_name_the_gateway_reach_the_provider_unchanged_and_never_the_peer( + gateway: Gateway, surface: Surface +) -> None: + with _rig(gateway, surface) as rig: + remote_mcp_tools: Final = [ + { + "type": "mcp", + "server_label": "docs", + "server_url": "https://mcp.example.com/sse", + "require_approval": "never", + }, + { + "type": "mcp", + "server_label": "repos", + "server_url": f"https://mcp.example.com/mcp/unserved{uuid.uuid4().hex[:8]}", + "require_approval": "never", + }, + ] + response: Final = rig.send(_granted_key(rig), [CALLER_FUNCTION_TOOLS[surface], *remote_mcp_tools]) + assert response.status_code == 200, response.text + bodies: Final = tuple( + JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST" + ) + assert len(bodies) == 1, bodies + forwarded: Final = bodies[0].get("tools") + assert isinstance(forwarded, list) and len(forwarded) == 3, forwarded + assert [tool for tool in forwarded if isinstance(tool, dict) and tool.get("type") == "mcp"] == remote_mcp_tools + assert "lookup" in json.dumps([tool for tool in forwarded if tool not in remote_mcp_tools]), forwarded + assert tool_calls(rig.peer.drain()) == (), "a tool that does not name the gateway reached the peer" + + def test_streaming_chat_executes_the_tool_once_and_streams_the_follow_up(gateway: Gateway) -> None: with _rig(gateway, "chat") as rig: key: Final = _granted_key(rig) diff --git a/tests/unit/responses/mcp/test_tool_precheck.py b/tests/unit/responses/mcp/test_tool_precheck.py new file mode 100644 index 00000000000..832151ce2a6 --- /dev/null +++ b/tests/unit/responses/mcp/test_tool_precheck.py @@ -0,0 +1,318 @@ +import json +import sys +from collections import deque +from collections.abc import Callable, Iterable, Mapping, Sequence +from importlib.abc import MetaPathFinder +from importlib.machinery import ModuleSpec +from types import ModuleType +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway + +GATEWAY_REGISTRY_MODULE: Final = "litellm.proxy._experimental.mcp_server.mcp_server_manager" +GATEWAY_HANDLER_MODULES: Final = ( + "litellm.responses.mcp.litellm_proxy_mcp_handler", + "litellm.responses.mcp.chat_completions_handler", + "litellm.llms.anthropic.pass_through.messages.mcp_handler", +) + +FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "noop", "description": "No operation", "parameters": {"type": "object", "properties": {}}}, +} +RESPONSES_FUNCTION_TOOL: Final = { + "type": "function", + "name": "noop", + "description": "No operation", + "parameters": {"type": "object", "properties": {}}, +} +ANTHROPIC_FUNCTION_TOOL: Final = { + "name": "noop", + "description": "No operation", + "input_schema": {"type": "object", "properties": {}}, +} +ANTHROPIC_MESSAGE_PAYLOAD: Final = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, +} +GATEWAY_MCP_TOOL: Final = {"type": "mcp", "server_label": "litellm", "server_url": "litellm_proxy"} +GATEWAY_SERVER_MCP_TOOL: Final = {"type": "mcp", "server_label": "x", "server_url": "litellm_proxy/mcp/github"} +GATEWAY_PATH_MCP_TOOL: Final = {"type": "mcp", "server_label": "x", "server_url": "https://gw.example.com/mcp/github"} +REMOTE_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/sse", + "require_approval": "never", +} + +RESPONSES_PAYLOAD: Final = { + "id": "resp_1", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "tools": [], +} + + +@pytest.fixture +def gateway_handler_unavailable(monkeypatch: pytest.MonkeyPatch) -> None: + for module_name in GATEWAY_HANDLER_MODULES: + monkeypatch.setitem(sys.modules, module_name, None) + + +class _FailingImport(MetaPathFinder): + def __init__(self, module_name: str, error: ImportError) -> None: + self.module_name: Final = module_name + self.error: Final = error + + def find_spec( + self, fullname: str, path: Sequence[str] | None, target: ModuleType | None = None + ) -> ModuleSpec | None: + if fullname == self.module_name: + raise self.error + return None + + +@pytest.fixture( + params=[ + ModuleNotFoundError("No module named 'mcp'"), + ImportError("cannot import name 'CacheableResult' from 'mcp.types'"), + ModuleNotFoundError("No module named 'fastapi'"), + ], + ids=["mcp-missing", "mcp-1x", "proxy-extras-missing"], +) +def gateway_registry_unavailable( + request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch, gateway_handler_unavailable: None +) -> None: + monkeypatch.delitem(sys.modules, GATEWAY_REGISTRY_MODULE, raising=False) + monkeypatch.setattr(sys, "meta_path", [_FailingImport(GATEWAY_REGISTRY_MODULE, request.param), *sys.meta_path]) + + +class _ToolList(list[Mapping[str, object]]): + pass + + +@pytest.mark.parametrize( + "tools", + [[FUNCTION_TOOL], [REMOTE_MCP_TOOL], _ToolList([FUNCTION_TOOL]), (FUNCTION_TOOL, REMOTE_MCP_TOOL)], + ids=["function-tool", "remote-mcp-tool", "list-subclass", "tuple"], +) +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_completion_without_a_gateway_tool_does_not_import_the_gateway_handler(tools: Iterable[Mapping[str, object]]): + response: Final = litellm.completion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + tools=tools, + mock_response="ok", + ) + + assert response.choices[0].message.content == "ok" + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("gateway_handler_unavailable") +async def test_acompletion_with_function_tool_does_not_import_the_gateway_handler(): + response: Final = await litellm.acompletion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + tools=[FUNCTION_TOOL], + mock_response="ok", + ) + + assert response.choices[0].message.content == "ok" + + +@pytest.mark.parametrize( + "tool", + [GATEWAY_MCP_TOOL, GATEWAY_SERVER_MCP_TOOL, GATEWAY_PATH_MCP_TOOL], + ids=["litellm-proxy", "litellm-proxy-server", "gateway-path-url"], +) +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_completion_with_gateway_mcp_tool_still_reaches_the_handler(tool: Mapping[str, object]): + with pytest.raises(ModuleNotFoundError, match=r"litellm_proxy_mcp_handler|chat_completions_handler"): + litellm.completion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + tools=[tool], + mock_response="ok", + ) + + +@pytest.mark.parametrize( + "tools", + [None, [RESPONSES_FUNCTION_TOOL], [REMOTE_MCP_TOOL]], + ids=["no-tools", "function-tool", "remote-mcp-tool"], +) +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_responses_without_a_gateway_tool_does_not_import_the_gateway_handler( + tools: list[Mapping[str, object]] | None, respx_mock: respx.MockRouter +): + upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=RESPONSES_PAYLOAD) + ) + + response: Final = litellm.responses(model="openai/gpt-4o-mini", input="hi", tools=tools, api_key="sk-test") + + assert upstream.called + assert response.output[0].content[0].text == "Done." + + +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_responses_with_gateway_mcp_tool_still_reaches_the_handler(): + with pytest.raises(litellm.APIConnectionError, match=r"litellm_proxy_mcp_handler"): + litellm.responses(model="openai/gpt-4o-mini", input="hi", tools=[GATEWAY_MCP_TOOL], api_key="sk-test") + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("gateway_handler_unavailable") +async def test_anthropic_messages_with_function_tool_does_not_import_the_gateway_handler( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=httpx.Response(200, json=ANTHROPIC_MESSAGE_PAYLOAD) + ) + + response: Final = await litellm.anthropic_interface.messages.acreate( + model="anthropic/claude-sonnet-5-5", + messages=[{"role": "user", "content": "hi"}], + max_tokens=16, + tools=[ANTHROPIC_FUNCTION_TOOL], + api_key="sk-test", + ) + + assert upstream.called + assert response["content"][0]["text"] == "ok" + + +@pytest.mark.parametrize( + "tools, expected", + [ + (None, False), + ([], False), + ([FUNCTION_TOOL], False), + ([RESPONSES_FUNCTION_TOOL], False), + ([{"type": "web_search"}], False), + (["not-a-mapping"], False), + ([REMOTE_MCP_TOOL], False), + ([{"type": "mcp", "server_label": "x"}], False), + ([{"type": "mcp", "server_label": "x", "server_url": None}], False), + ([{"server_url": "litellm_proxy"}], False), + ([FUNCTION_TOOL, GATEWAY_MCP_TOOL], True), + ([GATEWAY_SERVER_MCP_TOOL], True), + ([GATEWAY_PATH_MCP_TOOL], True), + ([REMOTE_MCP_TOOL, GATEWAY_PATH_MCP_TOOL], True), + ], + ids=[ + "none", + "empty", + "chat-function", + "responses-function", + "web-search", + "non-mapping", + "remote-mcp", + "mcp-without-url", + "mcp-null-url", + "gateway-url-without-mcp-type", + "gateway-among-functions", + "gateway-server", + "gateway-path-url", + "gateway-path-after-remote-mcp", + ], +) +def test_may_name_litellm_mcp_gateway_matches_only_gateway_mcp_tools(tools: list[object] | None, expected: bool): + assert may_name_litellm_mcp_gateway(tools) is expected + + +@pytest.mark.parametrize( + "tools, expected", + [ + ([GATEWAY_PATH_MCP_TOOL], False), + ([FUNCTION_TOOL, GATEWAY_PATH_MCP_TOOL], False), + ([GATEWAY_MCP_TOOL], True), + ([GATEWAY_SERVER_MCP_TOOL, GATEWAY_PATH_MCP_TOOL], True), + ], + ids=["gateway-path-url", "gateway-path-among-functions", "litellm-proxy", "litellm-proxy-server-and-path-url"], +) +@pytest.mark.usefixtures("gateway_registry_unavailable") +def test_proxy_path_url_names_the_gateway_only_when_its_registry_loads(tools: list[object], expected: bool): + assert may_name_litellm_mcp_gateway(tools) is expected + + +@pytest.mark.usefixtures("gateway_registry_unavailable") +def test_responses_sends_a_proxy_path_mcp_url_to_the_provider_without_the_gateway_registry( + respx_mock: respx.MockRouter, +): + upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=RESPONSES_PAYLOAD) + ) + + response: Final = litellm.responses( + model="openai/gpt-4o-mini", input="hi", tools=[GATEWAY_PATH_MCP_TOOL], api_key="sk-test" + ) + + assert response.output[0].content[0].text == "Done." + assert json.loads(upstream.calls.last.request.content)["tools"] == [GATEWAY_PATH_MCP_TOOL] + + +@pytest.mark.usefixtures("gateway_registry_unavailable") +def test_completion_sends_a_proxy_path_mcp_url_to_the_provider_without_the_gateway_registry(): + response: Final = litellm.completion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + tools=[GATEWAY_PATH_MCP_TOOL], + mock_response="ok", + ) + + assert response.choices[0].message.content == "ok" + + +class _ToolTuple(tuple[Mapping[str, object], ...]): + pass + + +@pytest.mark.parametrize( + "make_tools", + [ + lambda items: (item for item in items), + lambda items: iter(items), + lambda items: deque(items), + lambda items: _ToolList(items), + lambda items: _ToolTuple(items), + ], + ids=["generator", "iterator", "deque", "list-subclass", "tuple-subclass"], +) +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_precheck_reads_every_iterable_shape_without_the_gateway_handler( + make_tools: Callable[[list[Mapping[str, object]]], Iterable[Mapping[str, object]]], +): + assert may_name_litellm_mcp_gateway(make_tools([FUNCTION_TOOL, REMOTE_MCP_TOOL])) is False + assert may_name_litellm_mcp_gateway(make_tools([FUNCTION_TOOL, GATEWAY_MCP_TOOL])) is True diff --git a/tests/unit/responses/test_responses_prompt_management.py b/tests/unit/responses/test_responses_prompt_management.py index 1e8880424ea..9af92e7f9ca 100644 --- a/tests/unit/responses/test_responses_prompt_management.py +++ b/tests/unit/responses/test_responses_prompt_management.py @@ -70,10 +70,6 @@ def _patch_responses_dispatch(): import_module("litellm.responses.main").litellm, "get_llm_provider", side_effect=_provider_by_model, ), - patch.object( - import_module("litellm.responses.mcp.litellm_proxy_mcp_handler").LiteLLM_Proxy_MCP_Handler, "should_use_litellm_mcp_gateway", - return_value=False, - ), patch.object( import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config", return_value=None, @@ -158,7 +154,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm litellm.responses( @@ -193,7 +189,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm litellm.responses( @@ -216,7 +212,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm litellm.responses( @@ -247,7 +243,7 @@ class TestResponsesAPIPromptManagement: logging_obj.model_call_details = {} patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm litellm.responses( @@ -273,7 +269,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm litellm.responses( @@ -305,7 +301,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm litellm.responses( @@ -328,7 +324,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm litellm.responses( @@ -368,7 +364,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm litellm.responses( @@ -399,8 +395,7 @@ class TestResponsesAPIPromptManagement: side_effect=_provider_by_model, ), patches[1], - patches[2], - patches[3] as mock_handler, + patches[2] as mock_handler, ): import litellm @@ -437,7 +432,7 @@ class TestAsyncResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm await litellm.aresponses( @@ -467,7 +462,7 @@ class TestAsyncResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm await litellm.aresponses( @@ -498,7 +493,7 @@ class TestAsyncResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm await litellm.aresponses( @@ -524,7 +519,7 @@ class TestAsyncResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm await litellm.aresponses( @@ -657,7 +652,7 @@ async def test_async_guardrail_block_from_prompt_hook_reaches_caller_unwrapped() logging_obj.async_get_chat_completion_prompt = AsyncMock(side_effect=block) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3], pytest.raises(HTTPException) as exc_info: + with patches[0], patches[1], patches[2], pytest.raises(HTTPException) as exc_info: await litellm.aresponses(input="Hi", model="gpt-4o", prompt_id="blocked", litellm_logging_obj=logging_obj) assert exc_info.value is block @@ -669,7 +664,7 @@ def test_sync_guardrail_block_from_prompt_hook_reaches_caller_unwrapped(): logging_obj.get_chat_completion_prompt.side_effect = block patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3], pytest.raises(HTTPException) as exc_info: + with patches[0], patches[1], patches[2], pytest.raises(HTTPException) as exc_info: litellm.responses(input="Hi", model="gpt-4o", prompt_id="blocked", litellm_logging_obj=logging_obj) assert exc_info.value is block From ca4e6521e70b9ecdc13bf747b0c1fc48b174fcce Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 14:47:28 +0000 Subject: [PATCH 40/41] fix(router): serve a team's wildcard deployment before a proxy-wide wildcard (#48490) Backport of #48490 to lts/1.108.x. Cherry-picked from 76d001a7c06d5dbf93f9191b247db0ef981b6c66 (main). The line has no wildcard_routed_models helper and no get_models_by_pattern, so only wildcard_route_deployments is changed. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 12 +++---- .../test_ptu_consumption.py | 16 +++------ tests/unit/test_router/test_router.py | 36 ++++++++++++++++++- 3 files changed, 46 insertions(+), 18 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 244fac94d42..52b131362b1 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12915,13 +12915,13 @@ class Router: return [{**DeploymentTypedDict(**m), "model_name": model_name} for m in matches] def wildcard_route_deployments(self, model: str, team_id: str | None) -> Sequence[Mapping[str, object]]: - """The deployments routing serves ``model`` from through a wildcard: the most specific - proxy-wide pattern matching it, else ``team_id``'s own pattern, never both.""" - proxy_wide: Final[Sequence[Mapping[str, object]]] = self.pattern_router.get_deployments_by_pattern(model=model) + """The deployments routing serves ``model`` from through a wildcard: ``team_id``'s own pattern + matching it, else the most specific proxy-wide pattern, never both.""" team_router: Final = self.team_pattern_routers.get(team_id) if team_id is not None else None - if proxy_wide or team_router is None: - return proxy_wide - return team_router.get_deployments_by_pattern(model=model) + team_owned: Final[Sequence[Mapping[str, object]]] = ( + team_router.get_deployments_by_pattern(model=model) if team_router is not None else () + ) + return team_owned or self.pattern_router.get_deployments_by_pattern(model=model) def team_ptu_ceiling(self, team_id: str, model_group: str) -> PTUTeamCeiling | None: """The per-minute ceiling ``team_id``'s PTU shares set on the group serving ``model_group``, diff --git a/tests/unit/proxy/management_endpoints/test_ptu_consumption.py b/tests/unit/proxy/management_endpoints/test_ptu_consumption.py index a446afbd809..065ec93dcd0 100644 --- a/tests/unit/proxy/management_endpoints/test_ptu_consumption.py +++ b/tests/unit/proxy/management_endpoints/test_ptu_consumption.py @@ -232,14 +232,10 @@ def test_rows_keyed_by_a_name_a_wildcard_route_serves_are_sized_by_its_shared_de assert attached.results[0].breakdown.model_groups["ptu-chat"].metrics.ptu_hours == pytest.approx(1.0) -@pytest.mark.parametrize( - ("catch_all_routes", "expected_ptu_hours"), [(False, 1.0), (True, 0.0)], ids=["team-wildcard", "behind-catch-all"] -) -def test_a_teams_wildcard_reservation_reads_ptu_hours_only_while_the_router_serves_from_it( - monkeypatch, catch_all_routes, expected_ptu_hours -): - """The router tries the proxy-wide wildcards before a team's own, so a team's `ptu-*` - reservation behind an open `*` route never serves `ptu-chat`, and its row reads no PTU-hours.""" +@pytest.mark.parametrize("catch_all_routes", [False, True], ids=["team-wildcard", "beside-catch-all"]) +def test_a_teams_wildcard_reservation_reads_ptu_hours_even_beside_a_proxy_wide_catch_all(monkeypatch, catch_all_routes): + """The router tries a team's own wildcard before the proxy-wide ones, so a team's `ptu-*` + reservation serves `ptu-chat` whether or not an open `*` route exists, and its row reads its PTU-hours.""" monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") reservation: Final = { "model_name": "ptu-team-r", @@ -264,9 +260,7 @@ def test_a_teams_wildcard_reservation_reads_ptu_hours_only_while_the_router_serv attached: Final = with_ptu_consumption(_response(_day("2026-09-23", {"ptu-chat": one_hour})), router, "team-r") - assert attached.results[0].breakdown.model_groups["ptu-chat"].metrics.ptu_hours == pytest.approx( - expected_ptu_hours - ) + assert attached.results[0].breakdown.model_groups["ptu-chat"].metrics.ptu_hours == pytest.approx(1.0) def _mixed_ptu_router() -> Router: diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 8cde45e7603..fe9541cbe21 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -7371,7 +7371,8 @@ def test_get_wildcard_deployment_usable_by_team_prefers_the_team_pattern(): pytest.param(False, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-when-no-proxy-wide-match"), pytest.param(False, "rwild-chat", "team-b", (), id="another-teams-pattern-never-serves"), pytest.param(False, "rwild-chat", None, (), id="no-team-no-team-pattern"), - pytest.param(True, "rwild-chat", "team-r", ("catch-all",), id="proxy-wide-pattern-beats-the-team-pattern"), + pytest.param(True, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-beats-the-proxy-wide-pattern"), + pytest.param(True, "rwild-chat", "team-b", ("catch-all",), id="another-team-falls-through-to-proxy-wide"), pytest.param(True, "ptu-chat", "team-r", ("ptu-wildcard",), id="only-the-most-specific-proxy-wide-pattern"), ], ) @@ -7407,6 +7408,39 @@ def test_wildcard_route_deployments_is_the_one_pattern_routing_serves_from( assert tuple(Deployment.model_validate(d).model_info.id for d in early_deployments) == expected_ids +@pytest.mark.parametrize( + ("team_id", "expected_id"), + [ + pytest.param("team-a", "team-a-byok", id="team-key-uses-its-own-byok-wildcard"), + pytest.param("team-b", "shared-openrouter", id="team-without-byok-uses-the-shared-wildcard"), + pytest.param(None, "shared-openrouter", id="key-without-team-uses-the-shared-wildcard"), + ], +) +def test_team_byok_wildcard_is_served_before_a_proxy_wide_wildcard_of_the_same_name( + team_id: str | None, expected_id: str +): + """Regression: a proxy-wide openrouter/* shadowed a team's own openrouter/* BYOK deployment for that team's keys.""" + router: Final = litellm.Router( + model_list=[ + { + "model_name": "model_name_team-a_internal", + "litellm_params": {"model": "openrouter/*", "api_key": "sk-team-a"}, + "model_info": {"id": "team-a-byok", "team_id": "team-a", "team_public_model_name": "openrouter/*"}, + }, + { + "model_name": "openrouter/*", + "litellm_params": {"model": "openrouter/*", "api_key": "sk-shared"}, + "model_info": {"id": "shared-openrouter"}, + }, + ] + ) + requested: Final = "openrouter/openai/gpt-5-nano" + + served: Final = router.wildcard_route_deployments(requested, team_id) + + assert tuple(Deployment.model_validate(d).model_info.id for d in served) == (expected_id,) + + def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint(): """ Test that get_deployment_credentials_with_provider correctly copies From 854937381d0dd5ac5a7abf8421b7220d153eba12 Mon Sep 17 00:00:00 2001 From: ekpog200 <62500055+ekpog200@users.noreply.github.com> Date: Tue, 20 Oct 2026 08:42:59 +0300 Subject: [PATCH 41/41] fix(videos)!: keep model_id in proxy video ids so per-model keys resolve on status/content (#44614) Backport of #44614 to lts/1.108.x. Cherry-picked from ba61ecdf570c91a2409560138170ff9ab1fe7ad7 (main). The line has no splits_stream_timeout on _update_kwargs_with_deployment, so only the _video_call_pin wrapping is added around deployment selection. The line has no litellm/llms/base_llm/videos/async_task_utils.py, so its _id_safe_name dedupe is applied to the dashscope video transformation that holds the same helper on this line. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/dashscope/videos/transformation.py | 9 +- litellm/proxy/auth/auth_utils.py | 50 ++- litellm/proxy/video_endpoints/endpoints.py | 92 ++-- litellm/proxy/video_endpoints/utils.py | 115 ++++- litellm/router.py | 95 +++- litellm/router_utils/cooldown_handlers.py | 6 + .../router_utils/pattern_match_deployments.py | 2 +- litellm/types/videos/main.py | 2 + litellm/types/videos/utils.py | 111 +++-- .../router_code_coverage.py | 2 + .../routing/test_video_poll_routing_wire.py | 4 +- tests/unit/conftest.py | 10 +- tests/unit/proxy/auth/test_auth_utils.py | 412 ++++++++++++----- .../proxy/video_endpoints/test_endpoints.py | 162 ++++++- .../unit/proxy/video_endpoints/test_utils.py | 413 +++++++++++++++++- .../test_router_provider_endpoints.py | 247 +++++++++++ tests/unit/test_video_generation.py | 30 ++ 17 files changed, 1533 insertions(+), 229 deletions(-) diff --git a/litellm/llms/dashscope/videos/transformation.py b/litellm/llms/dashscope/videos/transformation.py index 17b685a0e1a..625ce3bab7e 100644 --- a/litellm/llms/dashscope/videos/transformation.py +++ b/litellm/llms/dashscope/videos/transformation.py @@ -35,6 +35,7 @@ from litellm.types.videos.utils import ( decode_video_id_with_provider, encode_video_id_with_provider, extract_original_video_id, + id_safe_name, ) from litellm.utils import get_model_info @@ -599,17 +600,13 @@ def _stamped_poll_models(metadata: object) -> tuple[str, str | None] | None: fields: Final = _PARAMETERS_ADAPTER.validate_python(metadata) except ValidationError: return None - requested: Final = _id_safe_name(fields.get("model_group")) - served: Final = _id_safe_name(fields.get("deployment_model_name")) + requested: Final = id_safe_name(fields.get("model_group")) + served: Final = id_safe_name(fields.get("deployment_model_name")) if requested is None: return None return requested, None if served is None or served == requested or "*" in served else served -def _id_safe_name(name: object) -> str | None: - return name if isinstance(name, str) and name and ";" not in name else None - - class DashScopeVideoConfig(BaseVideoConfig): def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseVideoConfig contract returns list return [ diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 98bb0985f4a..d012bbe51f3 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -4,6 +4,7 @@ import re import sys from collections.abc import Collection, Iterator, Mapping from functools import lru_cache +from itertools import accumulate, chain, repeat from logging import Logger from types import MappingProxyType from typing import Any, Final, Protocol @@ -39,6 +40,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment, server_owned_wif_fields_present from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params from litellm.types.utils import CustomPricingLiteLLMParams, oauth_token_exchange_litellm_params +from litellm.types.videos.main import DecodedVideoId def is_invalid_virtual_key_error(exception: BaseException | None) -> bool: @@ -566,6 +568,31 @@ def fallback_target_model_name(target: object) -> str | None: return None +def _direct_fallback_model_names(model: str, llm_router: Router) -> frozenset[str]: + from litellm.proxy.auth.model_checks import get_all_fallbacks + + return frozenset( + name for target in get_all_fallbacks(model, llm_router) if (name := fallback_target_model_name(target)) + ) + + +def _fallback_hop( + frontier_and_reached: tuple[frozenset[str], frozenset[str]], llm_router: Router +) -> tuple[frozenset[str], frozenset[str]]: + frontier, reached = frontier_and_reached + next_hop: Final = ( + frozenset(chain.from_iterable(_direct_fallback_model_names(model, llm_router) for model in frontier)) - reached + ) + return next_hop, reached | next_hop + + +def router_fallback_model_names(model: str, llm_router: Router) -> frozenset[str]: + """The groups a failure on `model` can reach through the router's fallbacks, chained and generic (`*`) ones.""" + hops: Final = accumulate(repeat(llm_router), _fallback_hop, initial=(frozenset({model}), frozenset({model}))) + _, reached = next(hop for hop in hops if not hop[0]) + return reached - {model} + + def request_fallback_model_names(request_body: Mapping[str, object]) -> tuple[str, ...]: return tuple( dict.fromkeys( @@ -1975,10 +2002,10 @@ def _extract_models_from_managed_resource_id( ) if resource_id_field == "video_id": - model_id = decode_video_id_with_provider(resource_id).get("model_id") + decoded_video: Final = decode_video_id_with_provider(resource_id) _append_model_candidates( candidates=candidates, - value=_resolve_model_id_with_router(model_id, llm_router), + value=_video_access_group(decoded_video, llm_router) or decoded_video.get("model_id"), ) else: model_id = decode_character_id_with_provider(resource_id).get("model_id") @@ -1992,6 +2019,25 @@ def _extract_models_from_managed_resource_id( return _dedupe_model_candidates(candidates) +def _video_access_group(decoded_video: DecodedVideoId, llm_router: Router | None) -> str | None: + """ + The group a key must be allowed to call to use a video: the creating deployment's group, or, while router + fallbacks skip model access checks, the group whose router fallback created it. + """ + from litellm.proxy.auth.fallback_model_access import router_fallback_access_check + from litellm.proxy.video_endpoints.utils import video_model_group + + creating_group: Final = video_model_group(decoded_video, llm_router) + requested_model: Final = decoded_video.get("requested_model") + if llm_router is None or creating_group is None or requested_model is None: + return creating_group + if router_fallback_access_check.is_enforced(): + return creating_group + if creating_group in router_fallback_model_names(requested_model, llm_router): + return requested_model + return creating_group + + def _resolve_model_id_with_router(model_id: str | None, llm_router: Router | None) -> str | None: if model_id is None or llm_router is None: return model_id diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index c4f33eae114..d4643994043 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -20,9 +20,14 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( ) from litellm.proxy.image_endpoints.endpoints import batch_to_bytesio from litellm.proxy.video_endpoints.utils import ( + VideoRequestRoute, encode_character_id_in_response, + encode_video_id_in_response, extract_model_from_target_model_names, + fallback_requested_model, get_custom_provider_from_data, + routable_model_name, + video_model_group, video_reference_to_id, ) from litellm.types.videos.utils import ( @@ -84,6 +89,7 @@ async def video_generation( # Read request body data: Final = await read_request_body(request=request) + requested: Final = VideoRequestRoute.model_validate(data) if input_reference is not None: input_reference_file: Final = await batch_to_bytesio([input_reference]) if input_reference_file: @@ -118,7 +124,13 @@ async def video_generation( version=version, ) else: - return generated + requested_model: Final = routable_model_name(requested.model, llm_router, requested.custom_llm_provider) + return encode_video_id_in_response( + response=generated, + fallback_provider=requested.custom_llm_provider, + fallback_model_id=requested_model, + requested_model=fallback_requested_model(requested_model, generated, llm_router), + ) @router.get( @@ -269,12 +281,9 @@ async def video_status( if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider - # Resolve model_name from model_id if available - # This allows the router to automatically inject litellm_params from the model config - if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id) - if resolved_model: - data["model"] = resolved_model + resolved_model: Final = video_model_group(decoded, llm_router) + if resolved_model: + data["model"] = resolved_model # Process request using ProxyBaseLLMRequestProcessing processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model")) @@ -305,7 +314,12 @@ async def video_status( version=version, ) else: - return status + return encode_video_id_in_response( + response=status, + fallback_provider=custom_llm_provider, + fallback_model_id=model_id_from_decoded, + polled_video=decoded, + ) @router.get( @@ -358,7 +372,6 @@ async def video_content( decoded: Final = decode_video_id_with_provider(video_id) provider_from_id: Final = decoded.get("custom_llm_provider") - model_id_from_decoded: Final = decoded.get("model_id") custom_llm_provider: Final = ( get_custom_llm_provider_from_request_headers(request=request) @@ -369,12 +382,9 @@ async def video_content( if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider - # Resolve model_name from model_id if available - # This allows the router to automatically inject litellm_params from the model config - if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id) - if resolved_model: - data["model"] = resolved_model + resolved_model: Final = video_model_group(decoded, llm_router) + if resolved_model: + data["model"] = resolved_model # Process request using ProxyBaseLLMRequestProcessing processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model")) try: @@ -462,6 +472,7 @@ async def video_remix( ) data: Final = await read_request_body(request=request) + requested: Final = VideoRequestRoute.model_validate(data) data["video_id"] = video_id decoded: Final = decode_video_id_with_provider(video_id) @@ -471,18 +482,15 @@ async def video_remix( custom_llm_provider: Final = ( get_custom_llm_provider_from_request_headers(request=request) or get_custom_llm_provider_from_request_query(request=request) - or data.get("custom_llm_provider") + or requested.custom_llm_provider or provider_from_id ) if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider - # Resolve model_name from model_id if available - # This allows the router to automatically inject litellm_params from the model config - if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id) - if resolved_model: - data["model"] = resolved_model + resolved_model: Final = video_model_group(decoded, llm_router) + if resolved_model: + data["model"] = resolved_model # Process request using ProxyBaseLLMRequestProcessing processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model")) @@ -513,7 +521,13 @@ async def video_remix( version=version, ) else: - return remixed + return encode_video_id_in_response( + response=remixed, + fallback_provider=custom_llm_provider, + fallback_model_id=model_id_from_decoded + or routable_model_name(requested.model, llm_router, custom_llm_provider), + requested_model=decoded.get("requested_model"), + ) @router.post( @@ -771,6 +785,7 @@ async def video_edit( ) data: Final = await read_request_body(request=request) + requested: Final = VideoRequestRoute.model_validate(data) uploaded_video: Final = data.pop("video", None) if isinstance(uploaded_video, StarletteUploadFile): video_files: Final = await batch_to_bytesio((uploaded_video,)) @@ -793,10 +808,9 @@ async def video_edit( ) data["custom_llm_provider"] = custom_llm_provider - if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id) - if resolved_model: - data["model"] = resolved_model + resolved_model: Final = video_model_group(decoded, llm_router) + if resolved_model: + data["model"] = resolved_model processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model")) try: @@ -826,7 +840,13 @@ async def video_edit( version=version, ) else: - return edited + return encode_video_id_in_response( + response=edited, + fallback_provider=custom_llm_provider, + fallback_model_id=model_id_from_decoded + or routable_model_name(requested.model, llm_router, custom_llm_provider), + requested_model=decoded.get("requested_model"), + ) @router.post( @@ -875,6 +895,7 @@ async def video_extension( ) data: Final = await read_request_body(request=request) + requested: Final = VideoRequestRoute.model_validate(data) data["video_id"] = video_reference_to_id(data.pop("video", None)) decoded: Final = decode_video_id_with_provider(data["video_id"]) @@ -890,10 +911,9 @@ async def video_extension( ) data["custom_llm_provider"] = custom_llm_provider - if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id) - if resolved_model: - data["model"] = resolved_model + resolved_model: Final = video_model_group(decoded, llm_router) + if resolved_model: + data["model"] = resolved_model processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model")) try: @@ -923,4 +943,10 @@ async def video_extension( version=version, ) else: - return extended + return encode_video_id_in_response( + response=extended, + fallback_provider=custom_llm_provider, + fallback_model_id=model_id_from_decoded + or routable_model_name(requested.model, llm_router, custom_llm_provider), + requested_model=decoded.get("requested_model"), + ) diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index 554f0f6ef53..a2be42f1b76 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -1,8 +1,18 @@ -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final import orjson +from pydantic import BaseModel, ConfigDict -from litellm.types.videos.utils import encode_character_id_with_provider +from litellm.types.videos.main import DecodedVideoId, VideoObject +from litellm.types.videos.utils import ( + decode_video_id_with_provider, + encode_character_id_with_provider, + encode_video_id_with_provider, + id_safe_name, +) + +if TYPE_CHECKING: + from litellm.router import Router def extract_model_from_target_model_names(target_model_names: Any) -> str | None: @@ -65,3 +75,104 @@ def encode_character_id_in_response(response: Any, custom_llm_provider: str, mod model_id=model_id, ) return response + + +class VideoRequestRoute(BaseModel): + model_config = ConfigDict(frozen=True) + + model: object = None + custom_llm_provider: object = None + + +def routable_model_name(model: object, llm_router: "Router | None", custom_llm_provider: object) -> str | None: + name: Final = id_safe_name(model) + if name is None or llm_router is None: + return name + provider: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else None + return name if llm_router.resolve_model_name_from_model_id(name, provider) else None + + +def video_model_group(decoded: DecodedVideoId, llm_router: "Router | None") -> str | None: + """ + The group a video id names. A provider model that several groups share resolves to the group of the + deployment that created the video, under its team's public name for a team model, while a name the router + routes by itself (the id's model, bare or provider-prefixed), or one that routes to that deployment's group + (an alias of it), is kept as written. + """ + if llm_router is None: + return None + model_id: Final = decoded.get("model_id") + provider: Final = decoded.get("custom_llm_provider") + resolved: Final = llm_router.resolve_model_name_from_model_id(model_id, provider) if model_id else None + deployment_id: Final = decoded.get("deployment_id") + if not deployment_id or resolved in (model_id, f"{provider}/{model_id}"): + return resolved + deployment: Final = llm_router.get_deployment(model_id=deployment_id) + if deployment is None: + return resolved + creating_group: Final = deployment.model_info.team_public_model_name or deployment.model_name + if resolved is not None and llm_router.resolve_served_model_name(resolved, creating_group) == creating_group: + return resolved + return creating_group + + +def fallback_requested_model(requested_model: str | None, response: object, llm_router: "Router | None") -> str | None: + """The group a create asked for when one of that group's router fallbacks served it.""" + from litellm.proxy.auth.auth_utils import router_fallback_model_names + + if requested_model is None or llm_router is None or not isinstance(response, VideoObject): + return None + deployment_id: Final = id_safe_name(response.hidden_params.get("model_id")) + deployment: Final = llm_router.get_deployment(model_id=deployment_id) if deployment_id else None + if deployment is None: + return None + creating_group: Final = deployment.model_info.team_public_model_name or deployment.model_name + return requested_model if creating_group in router_fallback_model_names(requested_model, llm_router) else None + + +def encode_video_id_in_response( + response: object, + fallback_provider: object, + fallback_model_id: object, + polled_video: DecodedVideoId | None = None, + requested_model: str | None = None, +) -> object: + """ + Stamps the provider, model and the deployment that answered into a returned video id. A poll of an existing + video (`polled_video`) answered without a routed deployment keeps the deployment its id already named, and + every returned id keeps the group a router fallback was serving (`requested_model`). + """ + if not isinstance(response, VideoObject) or not response.id: + return response + + decoded: Final = decode_video_id_with_provider(response.id) + hidden_params: Final = response.hidden_params + routed_deployment_id: Final = id_safe_name(hidden_params.get("model_id")) + deployment_id: Final = routed_deployment_id or (polled_video or decoded).get("deployment_id") + kept_requested_model: Final = ( + id_safe_name(requested_model) or decoded.get("requested_model") or (polled_video or {}).get("requested_model") + ) + model_id: Final = decoded.get("model_id") or id_safe_name(fallback_model_id) or routed_deployment_id + provider: Final = ( + id_safe_name(decoded.get("custom_llm_provider")) + or id_safe_name(hidden_params.get("custom_llm_provider")) + or id_safe_name(fallback_provider) + ) + if not model_id or not provider: + return response + if ( + model_id == decoded.get("model_id") + and deployment_id == decoded.get("deployment_id") + and kept_requested_model == decoded.get("requested_model") + ): + return response + + restamped_id: Final = encode_video_id_with_provider( + video_id=decoded.get("video_id") or response.id, + provider=provider, + model_id=model_id, + served_model=decoded.get("served_model"), + deployment_id=deployment_id, + requested_model=kept_requested_model, + ) + return response.model_copy(update={"id": restamped_id}) diff --git a/litellm/router.py b/litellm/router.py index 52b131362b1..9880ca6fdd0 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -34,6 +34,7 @@ from collections.abc import ( MutableMapping, Sequence, ) +from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache, partial @@ -223,6 +224,7 @@ from litellm.router_utils.cooldown_handlers import ( is_background_response_poll_not_found, is_caller_scoped_auth_failure, is_caller_timeout_408, + is_video_lookup_not_found, set_cooldown_deployments, ) from litellm.router_utils.deployment_pricing import deployment_pricing_models @@ -370,6 +372,7 @@ from litellm.types.utils import ( shared_backend_model_info, ) from litellm.types.utils import ModelInfo as ModelMapInfo +from litellm.types.videos.utils import VIDEO_ID_CALL_TYPES, decode_video_id_with_provider from litellm.utils import ( CustomStreamWrapper, EmbeddingResponse, @@ -488,6 +491,7 @@ def model_info_is_active_for_environment(model_info: Mapping[str, object] | None _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") _CallbackT = TypeVar("_CallbackT") +_DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) _ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset({"model", "api_base", "api_key", "api_version"}) _ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params" @@ -496,6 +500,7 @@ _CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600 CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_binding" _PRE_CALL_CHECK_FILTERED_PARAMS: Final = frozenset({"response_format"}) +_video_call_deployment_id: Final[ContextVar[str | None]] = ContextVar("video_call_deployment_id", default=None) _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType( { "prompt_caching": PromptCachingDeploymentCheck, @@ -5717,17 +5722,19 @@ class Router: passthrough_on_no_deployment: Final = kwargs.pop("passthrough_on_no_deployment", False) function_name: Final = "_ageneric_api_call_with_fallbacks" + handler_name: Final[str | None] = getattr(original_generic_function, "__name__", None) deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) try: - deployment = await self.async_get_available_deployment( # rebind-ok: set on success, see pre-init above - model=model, - request_kwargs=kwargs, - messages=kwargs.get("messages", None), - input=kwargs.get("input", None), - specific_deployment=kwargs.pop("specific_deployment", None), - ) + with self._video_call_pin(handler_name, kwargs): + deployment = await self.async_get_available_deployment( # rebind-ok: set on success, see pre-init + model=model, + request_kwargs=kwargs, + messages=kwargs.get("messages", None), + input=kwargs.get("input", None), + specific_deployment=kwargs.pop("specific_deployment", None), + ) except Exception as e: if passthrough_on_no_deployment: return await original_generic_function(model=model, **kwargs) @@ -6409,6 +6416,37 @@ class Router: ) return await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) + @staticmethod + @contextlib.contextmanager + def _video_call_pin(handler_name: str | None, request_kwargs: Mapping[str, object]) -> Generator[None]: + """ + Scopes the deployment a video id names to the routing of one call on that video (status, content, remix, + edit, extension). Any other call routes as if no id were sent, whatever its body carries. + """ + video_id: Final = request_kwargs.get("video_id") + deployment_id: Final = ( + decode_video_id_with_provider(video_id).get("deployment_id") + if handler_name in VIDEO_ID_CALL_TYPES and isinstance(video_id, str) + else None + ) + token: Final = _video_call_deployment_id.set(deployment_id) + try: + yield + finally: + _video_call_deployment_id.reset(token) + + @staticmethod + def _narrow_to_video_deployment(healthy_deployments: Sequence[_DeploymentT]) -> Sequence[_DeploymentT]: + """ + A video lives on the provider account of the deployment that created it, so a call on an id naming a + deployment that every routing filter kept goes to that deployment, else to whichever one routing picks. + """ + deployment_id: Final = _video_call_deployment_id.get() + creating: Final = tuple( + deployment for deployment in healthy_deployments if deployment_id in Router._deployment_ids((deployment,)) + ) + return creating or healthy_deployments + 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 @@ -6430,13 +6468,14 @@ class Router: kwargs=kwargs, metadata_variable_name=metadata_variable_name, ) - deployment: Final = self.get_available_deployment( - model=model, - messages=kwargs.get("messages", None), - input=kwargs.get("input", None), - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with self._video_call_pin(handler_name, kwargs): + deployment: Final = self.get_available_deployment( + model=model, + messages=kwargs.get("messages", None), + input=kwargs.get("input", None), + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="generic_api_call") data: Final = deployment["litellm_params"].copy() @@ -8914,6 +8953,13 @@ class Router: ) return False + if is_video_lookup_not_found(exception, kwargs): + verbose_router_logger.debug( + "Router: Exiting 'deployment_callback_on_failure' without cooldown. " + "Provider 404 says the video is not on this deployment's account, not that the deployment is down." + ) + return False + exception_status: Final = getattr(exception, "status_code", "") if is_caller_timeout_408(kwargs, exception_status): @@ -12337,7 +12383,8 @@ class Router: 1. First, check if model_id directly matches a model_name or deployment ID 2. If not, search through router's model_list to find a match by litellm_params.model 3. If not, keep model_id when the router routes that name itself (an alias, a routing group, or a wildcard - whose deployment serves `custom_llm_provider`) + whose deployment serves `custom_llm_provider`), else name it `custom_llm_provider/model_id` when only that + form routes to such a wildcard 4. Return the model_name if found, None otherwise Args: @@ -12385,8 +12432,15 @@ class Router: if self.is_recognized_model(model_id): return model_id - if custom_llm_provider is not None and self.pattern_router.routes_to_provider(model_id, custom_llm_provider): + if custom_llm_provider is None: + return None + if self.pattern_router.routes_to_provider(model_id, custom_llm_provider): return model_id + provider_prefixed: Final = f"{custom_llm_provider}/{model_id}" + if not model_id.startswith(f"{custom_llm_provider}/") and self.pattern_router.routes_to_provider( + provider_prefixed, custom_llm_provider + ): + return provider_prefixed return None def resolve_served_model_name(self, model_name: str, served_model: str) -> str: @@ -13897,10 +13951,10 @@ class Router: def _filter_deployments_by_model_access_groups( self, model: str, - healthy_deployments: list, + healthy_deployments: list[DeploymentTypedDict], request_kwargs: Mapping[str, object] | None, request_team_id: str | None, - ) -> list: + ) -> list[DeploymentTypedDict]: """ Restrict candidate deployments to caller-authorized model access groups. @@ -14084,7 +14138,8 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) healthy_deployments = litellm.utils.get_order_filtered_deployments( - cast(list[dict], healthy_deployments), target_order=_target_order + list(self._narrow_to_video_deployment(cast(list[dict[str, object]], healthy_deployments))), + target_order=_target_order, ) ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in @@ -15271,7 +15326,7 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) healthy_deployments = litellm.utils.get_order_filtered_deployments( - healthy_deployments, target_order=_target_order + list(self._narrow_to_video_deployment(healthy_deployments)), target_order=_target_order ) ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index cdca5009b0e..20ad8a6e163 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -35,6 +35,7 @@ from litellm.types.utils import ( BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN, InternalCallOrigin, ) +from litellm.types.videos.utils import VIDEO_ID_CALL_TYPES from .router_callbacks.track_deployment_metrics import ( get_deployment_failures_for_current_minute, @@ -89,6 +90,11 @@ def is_background_response_poll_not_found(exception: Exception, litellm_params: ) +def is_video_lookup_not_found(exception: object, failure_kwargs: Mapping[str, object]) -> bool: + """Whether a call on an existing video id failed with a provider 404: the video is not on that account.""" + return getattr(exception, "status_code", None) == 404 and failure_kwargs.get("call_type") in VIDEO_ID_CALL_TYPES + + _EXCEPTION_POLICY_FIELDS: Final[tuple[tuple[type, str], ...]] = ( # ContentPolicyViolationError subclasses BadRequestError, so it must be checked first. (litellm.ContentPolicyViolationError, "ContentPolicyViolationErrorAllowedFails"), diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py index 62ba5592f90..daa052a46e1 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/router_utils/pattern_match_deployments.py @@ -173,7 +173,7 @@ class PatternMatchRouter: return None # No matching pattern found def routes_to_provider(self, request: str, custom_llm_provider: str) -> bool: - deployments: Final = _ROUTED_DEPLOYMENTS.validate_python(self.route(request) or ()) + deployments: Final = _ROUTED_DEPLOYMENTS.validate_python(self.get_pattern(request) or ()) return any( (deployment.litellm_params.custom_llm_provider or inferred_provider(deployment.litellm_params.model)) == custom_llm_provider diff --git a/litellm/types/videos/main.py b/litellm/types/videos/main.py index 05444b1e7ec..4d6f8fc6a30 100644 --- a/litellm/types/videos/main.py +++ b/litellm/types/videos/main.py @@ -113,6 +113,8 @@ class DecodedVideoId(TypedDict, total=False): model_id: str | None video_id: str served_model: ReadOnly[str | None] + deployment_id: ReadOnly[str | None] + requested_model: ReadOnly[str | None] class CharacterObject(LiteLLMBaseModel): diff --git a/litellm/types/videos/utils.py b/litellm/types/videos/utils.py index 9a07c9ba802..0fc2b8ede1b 100644 --- a/litellm/types/videos/utils.py +++ b/litellm/types/videos/utils.py @@ -6,6 +6,7 @@ Format: vid_{base64_encoded_string} """ import base64 +from collections.abc import Sequence from typing import Final from litellm._logging import verbose_logger @@ -14,6 +15,22 @@ from litellm.types.videos.main import DecodedVideoId VIDEO_ID_PREFIX: Final = "video_" VIDEO_ID_SERVED_MODEL_PREFIX: Final = "served_model:" +VIDEO_ID_DEPLOYMENT_PREFIX: Final = "deployment_id:" +VIDEO_ID_REQUESTED_MODEL_PREFIX: Final = "requested_model:" +VIDEO_ID_CALL_TYPES: Final = frozenset( + { + "avideo_status", + "video_status", + "avideo_content", + "video_content", + "avideo_remix", + "video_remix", + "avideo_edit", + "video_edit", + "avideo_extension", + "video_extension", + } +) CHARACTER_ID_PREFIX: Final = "character_" CHARACTER_ID_TEMPLATE: Final = "litellm:custom_llm_provider:{};model_id:{};character_id:{}" @@ -37,7 +54,12 @@ def _add_base64_padding(value: str) -> str: def encode_video_id_with_provider( - video_id: str, provider: str, model_id: str | None = None, served_model: str | None = None + video_id: str, + provider: str, + model_id: str | None = None, + served_model: str | None = None, + deployment_id: str | None = None, + requested_model: str | None = None, ) -> str: """Encode provider and model_id into video_id using base64.""" if not provider or not video_id: @@ -54,12 +76,24 @@ def encode_video_id_with_provider( # ID is not encoded (even if it starts with video_), so encode it assembled_id = str(SpecialEnums.LITELLM_MANAGED_VIDEO_COMPLETE_STR.value).format(provider, model_id or "", video_id) served_part: Final = f";{VIDEO_ID_SERVED_MODEL_PREFIX}{served_model}" if served_model else "" + deployment_part: Final = f";{VIDEO_ID_DEPLOYMENT_PREFIX}{deployment_id}" if deployment_id else "" + requested_part: Final = f";{VIDEO_ID_REQUESTED_MODEL_PREFIX}{requested_model}" if requested_model else "" - base64_encoded_id: Final[str] = base64.b64encode((assembled_id + served_part).encode("utf-8")).decode("utf-8") + base64_encoded_id: Final[str] = base64.b64encode( + (assembled_id + served_part + deployment_part + requested_part).encode("utf-8") + ).decode("utf-8") return f"{VIDEO_ID_PREFIX}{base64_encoded_id}" +def id_safe_name(name: object) -> str | None: + return name if isinstance(name, str) and name and ";" not in name else None + + +def _trailing_part(parts: Sequence[str], prefix: str) -> str | None: + return next((part.removeprefix(prefix) or None for part in parts[3:] if part.startswith(prefix)), None) + + def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId: """Decode provider and model_id from encoded video_id.""" if not encoded_video_id: @@ -80,43 +114,6 @@ def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId: cleaned_id = encoded_video_id.replace(VIDEO_ID_PREFIX, "") cleaned_id = _add_base64_padding(cleaned_id) decoded_id: Final = base64.b64decode(cleaned_id.encode("utf-8")).decode("utf-8") - - if ";" not in decoded_id: - return DecodedVideoId( - custom_llm_provider=None, - model_id=None, - video_id=encoded_video_id, - ) - - parts: Final = decoded_id.split(";") - - custom_llm_provider = None - model_id = None - decoded_video_id = encoded_video_id - - if len(parts) >= 3: - custom_llm_provider_part: Final = parts[0] - model_id_part: Final = parts[1] - video_id_part: Final = parts[2] - - custom_llm_provider = custom_llm_provider_part.replace("litellm:custom_llm_provider:", "") - model_id = model_id_part.replace("model_id:", "") - decoded_video_id = video_id_part.replace("video_id:", "") - - served_model: Final = ( - parts[3].removeprefix(VIDEO_ID_SERVED_MODEL_PREFIX) or None - if len(parts) >= 4 and parts[3].startswith(VIDEO_ID_SERVED_MODEL_PREFIX) - else None - ) - - if served_model is None: - return DecodedVideoId(custom_llm_provider=custom_llm_provider, model_id=model_id, video_id=decoded_video_id) - return DecodedVideoId( - custom_llm_provider=custom_llm_provider, - model_id=model_id, - video_id=decoded_video_id, - served_model=served_model, - ) except Exception as e: verbose_logger.debug("Error decoding video_id '%s': %s", encoded_video_id, e) return DecodedVideoId( @@ -125,6 +122,42 @@ def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId: video_id=encoded_video_id, ) + if ";" not in decoded_id: + return DecodedVideoId( + custom_llm_provider=None, + model_id=None, + video_id=encoded_video_id, + ) + + parts: Final = decoded_id.split(";") + + custom_llm_provider = None + model_id = None + decoded_video_id = encoded_video_id + + if len(parts) >= 3: + custom_llm_provider_part: Final = parts[0] + model_id_part: Final = parts[1] + video_id_part: Final = parts[2] + + custom_llm_provider = custom_llm_provider_part.replace("litellm:custom_llm_provider:", "") + model_id = model_id_part.replace("model_id:", "") + decoded_video_id = video_id_part.replace("video_id:", "") + + served_model: Final = _trailing_part(parts, VIDEO_ID_SERVED_MODEL_PREFIX) + deployment_id: Final = _trailing_part(parts, VIDEO_ID_DEPLOYMENT_PREFIX) + requested_model: Final = _trailing_part(parts, VIDEO_ID_REQUESTED_MODEL_PREFIX) + decoded_parts: Final = DecodedVideoId( + custom_llm_provider=custom_llm_provider, model_id=model_id, video_id=decoded_video_id + ) + with_served: Final[DecodedVideoId] = ( + decoded_parts if served_model is None else {**decoded_parts, "served_model": served_model} + ) + with_deployment: Final[DecodedVideoId] = ( + with_served if deployment_id is None else {**with_served, "deployment_id": deployment_id} + ) + return with_deployment if requested_model is None else {**with_deployment, "requested_model": requested_model} + def extract_original_video_id(encoded_video_id: str) -> str: """Extract original video ID without encoding.""" diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 9bd077ef3a5..8e84cf2e3e9 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -99,6 +99,8 @@ ignored_function_names = [ "_embedding", "_aembedding", "_stream_pre_content_error", # Tested through the non-retriable retry error tests in test_router.py + "_narrow_to_video_deployment", # Tested through video_status and avideo_status polls in test_router_provider_endpoints.py + "_video_call_pin", # Tested through the video_status pin tests and the chat video_id test in test_router_provider_endpoints.py "_deployment_num_retries", # Tested through the deployment num_retries mid-stream budget test in test_router.py "_request_fallback_list", # Tested through every mid-stream retry test in test_router.py "_request_model_group", # Tested through test_anthropic_messages_retry_budget_precedence_direct_call diff --git a/tests/integration/routing/test_video_poll_routing_wire.py b/tests/integration/routing/test_video_poll_routing_wire.py index e0d88a6b0ff..a3d767275a0 100644 --- a/tests/integration/routing/test_video_poll_routing_wire.py +++ b/tests/integration/routing/test_video_poll_routing_wire.py @@ -338,9 +338,9 @@ def test_an_id_without_a_served_group_still_polls_its_own_deployment(rig: _Rig) video_id: Final = created_id(create(rig.gateway, {"model": rig.member_a, "seconds": "5", "size": "1280x720"})) assert decoded(video_id).get("served_model") is None legacy: Final = forged_id("dashscope", rig.member_a, task_of(video_id), None) - assert legacy == video_id + assert {**decoded(legacy), "deployment_id": decoded(video_id).get("deployment_id")} == decoded(video_id) rig.drain() - _assert_served(_poll(rig, legacy), legacy) + _assert_served(_poll(rig, legacy), video_id) download: Final = _download(rig, legacy) assert download.status_code == 200 and download.content == MP4, download.text received: Final = rig.a.drain() diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 8331fb1937b..9028b8461b0 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -93,7 +93,15 @@ CALLBACK_LISTS: Final = ( "_async_failure_callback", "_async_input_callback", ) -RESET_TO_NONE_GLOBALS: Final = ("model_fallbacks", "cache") +RESET_TO_NONE_GLOBALS: Final = ( + "model_fallbacks", + "cache", + "max_fallbacks", + "default_fallbacks", + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", +) RESTORED_GLOBALS: Final = ( "disable_aiohttp_transport", "force_ipv4", diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index eb76901db45..a8b2ff4f6cd 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -414,18 +414,11 @@ class TestLogOnceIfBudgetReservationDisabled: def test_logs_info_only_once_when_enabled(self, caplog): with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): log_once_if_budget_reservation_disabled(disabled=False) - assert not any( - "disable_budget_reservation is enabled" in record.message - for record in caplog.records - ) + assert not any("disable_budget_reservation is enabled" in record.message for record in caplog.records) for _ in range(3): log_once_if_budget_reservation_disabled(disabled=True) - records = [ - record - for record in caplog.records - if "disable_budget_reservation is enabled" in record.message - ] + records = [record for record in caplog.records if "disable_budget_reservation is enabled" in record.message] assert len(records) == 1 assert records[0].levelno == logging.INFO @@ -753,19 +746,28 @@ def test_get_model_from_request_no_request_extracts_model(): ) -@pytest.mark.parametrize("provider,model", [ - ("laya", "english"), ("laya", "multilingual"), ("laya", "typed-decisions"), - ("bespoke", "nimble-latest"), ("bespoke", "bespokelabs/Bespoke-Nimble-9B"), -]) +@pytest.mark.parametrize( + "provider,model", + [ + ("laya", "english"), + ("laya", "multilingual"), + ("laya", "typed-decisions"), + ("bespoke", "nimble-latest"), + ("bespoke", "bespokelabs/Bespoke-Nimble-9B"), + ], +) @pytest.mark.parametrize("suffix", ["", "/"]) def test_oss_native_model_uses_the_classifier_permission_identity(provider: str, model: str, suffix: str) -> None: - assert get_model_from_request( - request_data={"model": model}, route=f"/{provider}/v1/systemone{suffix}" - ) == f"{provider}/{model}" + assert ( + get_model_from_request(request_data={"model": model}, route=f"/{provider}/v1/systemone{suffix}") + == f"{provider}/{model}" + ) @pytest.mark.parametrize("provider", ["laya", "bespoke"]) -@pytest.mark.parametrize("model", [None, "", "auto", "laya/english", "bespoke/nimble-latest", "unknown", ["english"], 7]) +@pytest.mark.parametrize( + "model", [None, "", "auto", "laya/english", "bespoke/nimble-latest", "unknown", ["english"], 7] +) def test_oss_native_model_cannot_implicitly_select_an_unauthorized_checkpoint(provider: str, model: object) -> None: with pytest.raises(HTTPException) as denied: get_model_from_request(request_data={"model": model}, route=f"/{provider}/v1/systemone") @@ -780,38 +782,49 @@ def test_laya_model_normalization_does_not_change_other_provider_routes() -> Non def _cache_prediction_router(): from litellm.router import Router - return Router(model_list=[ - { - "model_name": group, - "litellm_params": {"model": "anthropic/claude-sonnet-5", "api_key": "test-provider-key"}, - "model_info": {"id": deployment_id, "team_id": team_id}, - } - for group, deployment_id, team_id in ( - ("current-group", "current-id", None), ("candidate-group", "candidate-id", None), - ("own-group", "own-id", "prediction-team"), ("foreign-group", "foreign-id", "foreign-team"), - ) - ]) + return Router( + model_list=[ + { + "model_name": group, + "litellm_params": {"model": "anthropic/claude-sonnet-5", "api_key": "test-provider-key"}, + "model_info": {"id": deployment_id, "team_id": team_id}, + } + for group, deployment_id, team_id in ( + ("current-group", "current-id", None), + ("candidate-group", "candidate-id", None), + ("own-group", "own-id", "prediction-team"), + ("foreign-group", "foreign-id", "foreign-team"), + ) + ] + ) -@pytest.mark.parametrize("candidate,team_id,expected", [ - ("candidate-id", None, ["current-group", "candidate-group"]), - ("current-id", None, "current-group"), - ("missing-id", None, None), - ("candidate-group", None, None), - ("own-id", None, None), - ("own-id", "prediction-team", ["current-group", "own-group"]), - ("foreign-id", "prediction-team", None), -]) +@pytest.mark.parametrize( + "candidate,team_id,expected", + [ + ("candidate-id", None, ["current-group", "candidate-group"]), + ("current-id", None, "current-group"), + ("missing-id", None, None), + ("candidate-group", None, None), + ("own-id", None, None), + ("own-id", "prediction-team", ["current-group", "own-group"]), + ("foreign-id", "prediction-team", None), + ], +) def test_cache_prediction_auth_resolves_only_exact_deployment_ids(candidate, team_id, expected): - assert get_model_from_request( - request_data={ - "current_deployment_id": "current-id", "candidate_deployment_id": candidate, - "request": {"model": "caller-controlled-provider-model"}, - }, - route="/cost/predict-cache", - llm_router=_cache_prediction_router(), - team_id=team_id, - ) == expected + assert ( + get_model_from_request( + request_data={ + "current_deployment_id": "current-id", + "candidate_deployment_id": candidate, + "request": {"model": "caller-controlled-provider-model"}, + }, + route="/cost/predict-cache", + llm_router=_cache_prediction_router(), + team_id=team_id, + ) + == expected + ) def _cache_prediction_auth_app( @@ -834,18 +847,29 @@ def _cache_prediction_auth_app( router = _cache_prediction_router() allowed_models = ["current-group", "candidate-group", "own-group"] token = UserAPIKeyAuth( - api_key="test-proxy-key-hash", user_id="prediction-user", user_role=LitellmUserRoles.INTERNAL_USER, - models=allowed_models if key_models is None else key_models, team_id=team_id, + api_key="test-proxy-key-hash", + user_id="prediction-user", + user_role=LitellmUserRoles.INTERNAL_USER, + models=allowed_models if key_models is None else key_models, + team_id=team_id, team_models=allowed_models if team_models is None else team_models, - allowed_routes=allowed_routes, metadata=metadata or {}, + allowed_routes=allowed_routes, + metadata=metadata or {}, ) user = LiteLLM_UserTable( - user_id=token.user_id, user_role=LitellmUserRoles.INTERNAL_USER.value, models=user_models, + user_id=token.user_id, + user_role=LitellmUserRoles.INTERNAL_USER.value, + models=user_models, ) + async def authenticate(request, request_data, **_headers): await auth._enforce_key_and_fallback_model_access( - valid_token=token, request_data=request_data, route=request.url.path, request=request, - llm_model_list=router.get_model_list(), llm_router=router, + valid_token=token, + request_data=request_data, + route=request.url.path, + request=request, + llm_model_list=router.get_model_list(), + llm_router=router, ) return token @@ -878,22 +902,36 @@ def _cache_prediction_auth_app( def _cache_prediction_payload(candidate="candidate-id", current="current-id"): return { - "current_deployment_id": current, "candidate_deployment_id": candidate, - "request": {"messages": [{"role": "user", "content": [{ - "type": "text", "text": "Stable cached context", - "cache_control": {"type": "ephemeral"}, - }]}]}, + "current_deployment_id": current, + "candidate_deployment_id": candidate, + "request": { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Stable cached context", + "cache_control": {"type": "ephemeral"}, + } + ], + } + ] + }, } @pytest.mark.asyncio -@pytest.mark.parametrize("allowed_routes,user_models,candidate,status_code", [ - (["/chat/completions"], ["current-group", "candidate-group"], "candidate-id", 403), - (["/cost/predict-cache"], ["current-group"], "candidate-id", 403), - (["/cost/*"], ["current-group", "candidate-group"], "candidate-id", 200), - (["/cost/predict-cache"], ["current-group"], "current-id", 200), - (["/cost/predict-cache"], ["current-group"], "missing-id", 404), -]) +@pytest.mark.parametrize( + "allowed_routes,user_models,candidate,status_code", + [ + (["/chat/completions"], ["current-group", "candidate-group"], "candidate-id", 403), + (["/cost/predict-cache"], ["current-group"], "candidate-id", 403), + (["/cost/*"], ["current-group", "candidate-group"], "candidate-id", 200), + (["/cost/predict-cache"], ["current-group"], "current-id", 200), + (["/cost/predict-cache"], ["current-group"], "missing-id", 404), + ], +) async def test_cache_prediction_authorizes_route_and_personal_models_before_provider_counts( monkeypatch, allowed_routes, user_models, candidate, status_code ): @@ -912,20 +950,27 @@ async def test_cache_prediction_authorizes_route_and_personal_models_before_prov @pytest.mark.asyncio @pytest.mark.parametrize("arm", ["current_deployment_id", "candidate_deployment_id"]) -@pytest.mark.parametrize("team_id,key_models,user_models,team_models", [ - (None, ["*"], ["*"], None), - (None, ["current-group", "candidate-group"], ["*"], None), - (None, ["*"], ["current-group", "candidate-group"], None), - ("prediction-team", ["*"], ["*"], ["current-group", "candidate-group"]), -]) +@pytest.mark.parametrize( + "team_id,key_models,user_models,team_models", + [ + (None, ["*"], ["*"], None), + (None, ["current-group", "candidate-group"], ["*"], None), + (None, ["*"], ["current-group", "candidate-group"], None), + ("prediction-team", ["*"], ["*"], ["current-group", "candidate-group"]), + ], +) async def test_cache_prediction_hides_foreign_and_missing_ids_before_model_authorization( monkeypatch, arm, team_id, key_models, user_models, team_models ): import httpx app, counts = _cache_prediction_auth_app( - monkeypatch, ["/cost/predict-cache"], user_models, - team_id=team_id, key_models=key_models, team_models=team_models, + monkeypatch, + ["/cost/predict-cache"], + user_models, + team_id=team_id, + key_models=key_models, + team_models=team_models, ) async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: missing = await client.post("/cost/predict-cache", json={**_cache_prediction_payload(), arm: "missing-id"}) @@ -938,19 +983,26 @@ async def test_cache_prediction_hides_foreign_and_missing_ids_before_model_autho @pytest.mark.asyncio @pytest.mark.parametrize("arm", ["current_deployment_id", "candidate_deployment_id"]) -@pytest.mark.parametrize("key_models,team_models,status_code", [ - (["*"], ["*"], 200), - (["current-group", "candidate-group"], ["*"], 403), - (["*"], ["current-group", "candidate-group"], 403), -]) +@pytest.mark.parametrize( + "key_models,team_models,status_code", + [ + (["*"], ["*"], 200), + (["current-group", "candidate-group"], ["*"], 403), + (["*"], ["current-group", "candidate-group"], 403), + ], +) async def test_cache_prediction_checks_each_visible_team_deployment_model( monkeypatch, arm, key_models, team_models, status_code ): import httpx app, counts = _cache_prediction_auth_app( - monkeypatch, ["/cost/predict-cache"], ["*"], - team_id="prediction-team", key_models=key_models, team_models=team_models, + monkeypatch, + ["/cost/predict-cache"], + ["*"], + team_id="prediction-team", + key_models=key_models, + team_models=team_models, ) async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: response = await client.post("/cost/predict-cache", json={**_cache_prediction_payload(), arm: "own-id"}) @@ -976,19 +1028,24 @@ async def test_cache_prediction_checks_each_visible_personal_deployment_model(mo @pytest.mark.asyncio -@pytest.mark.parametrize("header_tag,key_tags,limit,status_code,provider_calls", [ - ("limited", [], 1, 429, 1), - (None, ["limited"], 1, 429, 1), - ("limited", ["limited"], 4, 200, 4), - ("unlimited", [], 1, 200, 4), -]) +@pytest.mark.parametrize( + "header_tag,key_tags,limit,status_code,provider_calls", + [ + ("limited", [], 1, 429, 1), + (None, ["limited"], 1, 429, 1), + ("limited", ["limited"], 4, 200, 4), + ("unlimited", [], 1, 200, 4), + ], +) async def test_cache_prediction_preserves_authenticated_header_and_key_tag_rpm( monkeypatch, header_tag, key_tags, limit, status_code, provider_calls ): import httpx app, counts = _cache_prediction_auth_app( - monkeypatch, ["/cost/predict-cache"], ["current-group", "candidate-group"], + monkeypatch, + ["/cost/predict-cache"], + ["current-group", "candidate-group"], metadata={"tag_rpm_limit": {"limited": limit}, "tags": key_tags}, ) headers = {"x-litellm-tags": header_tag} if header_tag else {} @@ -1142,7 +1199,11 @@ def _azure_relay_router(): model_list=[ { "model_name": "gpt", - "litellm_params": {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://a.services.ai.azure.com", "api_key": "k"}, + "litellm_params": { + "model": "azure_ai/gpt-5.4-mini", + "api_base": "https://a.services.ai.azure.com", + "api_key": "k", + }, }, { "model_name": "other-group", @@ -1364,8 +1425,14 @@ def test_get_model_from_request_resolves_video_id_model_with_router(): provider="vertex_ai", model_id="veo-3.1-generate-001", ) - llm_router = MagicMock() - llm_router.resolve_model_name_from_model_id.return_value = "gcp/google/veo-3.1-generate-001" + llm_router = Router( + model_list=[ + { + "model_name": "gcp/google/veo-3.1-generate-001", + "litellm_params": {"model": "vertex_ai/veo-3.1-generate-001"}, + } + ] + ) assert ( get_model_from_request( @@ -1375,7 +1442,148 @@ def test_get_model_from_request_resolves_video_id_model_with_router(): ) == "gcp/google/veo-3.1-generate-001" ) - llm_router.resolve_model_name_from_model_id.assert_called_once_with("veo-3.1-generate-001") + + +def test_get_model_from_request_checks_the_group_of_the_deployment_that_created_the_video() -> None: + from litellm.types.videos.utils import encode_video_id_with_provider + + llm_router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": "azure/sora-2", "api_base": f"https://{name}.example.com", "api_key": name}, + "model_info": {"id": f"dep-{name}"}, + } + for name in ("sora-team", "sora-shared") + ] + ) + video_id: Final = encode_video_id_with_provider( + video_id="video_raw", provider="azure", model_id="sora-2", deployment_id="dep-sora-shared" + ) + + assert ( + get_model_from_request( + request_data={"video_id": video_id}, + route="/v1/videos/{video_id}", + llm_router=llm_router, + ) + == "sora-shared" + ) + + +def test_get_model_from_request_checks_a_video_created_through_a_provider_wildcard_by_its_own_name() -> None: + from litellm.types.videos.utils import encode_video_id_with_provider + + llm_router: Final = Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "wildcard-key"}, + "model_info": {"id": "dep-wildcard"}, + } + ] + ) + video_id: Final = encode_video_id_with_provider( + video_id="video_raw", provider="openai", model_id="sora-2", deployment_id="dep-wildcard" + ) + routed_ids: Final = [ + deployment["model_info"]["id"] for deployment in llm_router.pattern_router.get_deployments_by_pattern("sora-2") + ] + + assert routed_ids == ["dep-wildcard"] + assert ( + get_model_from_request( + request_data={"video_id": video_id}, + route="/v1/videos/{video_id}", + llm_router=llm_router, + ) + == "sora-2" + ) + + +def test_get_model_from_request_checks_a_wildcard_video_whose_bare_name_does_not_route_by_its_prefixed_name() -> None: + from litellm.types.videos.utils import encode_video_id_with_provider + + llm_router: Final = Router( + model_list=[ + { + "model_name": "gemini/*", + "litellm_params": {"model": "gemini/*", "api_key": "wildcard-key"}, + "model_info": {"id": "dep-wildcard"}, + } + ] + ) + video_id: Final = encode_video_id_with_provider( + video_id="operations/op_1", provider="gemini", model_id="unlisted-video-model", deployment_id="dep-wildcard" + ) + + assert ( + get_model_from_request( + request_data={"video_id": video_id}, + route="/v1/videos/{video_id}", + llm_router=llm_router, + ) + == "gemini/unlisted-video-model" + ) + + +def _sora_fallback_router() -> Router: + return Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": "azure/sora-2", "api_base": f"https://{name}.example.com", "api_key": name}, + "model_info": {"id": f"dep-{name}"}, + } + for name in ("sora-pair", "sora-backup", "sora-last", "sora-other") + ], + fallbacks=[{"sora-pair": ["sora-backup"]}, {"sora-backup": ["sora-last"]}], + ) + + +@pytest.mark.parametrize( + ("general_settings", "requested_model", "creating_group", "expected"), + [ + ({"enforce_fallback_model_access": False}, "sora-pair", "sora-backup", "sora-pair"), + ({"enforce_fallback_model_access": False}, "sora-pair", "sora-last", "sora-pair"), + ({}, "sora-pair", "sora-backup", "sora-backup"), + ({"enforce_fallback_model_access": False}, "sora-other", "sora-backup", "sora-backup"), + ({"enforce_fallback_model_access": False}, None, "sora-backup", "sora-backup"), + ], + ids=[ + "opt-out-fallback", + "opt-out-chained-fallback", + "enforced", + "opt-out-no-fallback-relation", + "opt-out-no-requested-model", + ], +) +def test_get_model_from_request_checks_a_fallback_served_video_by_the_requested_group_only_under_the_opt_out( + monkeypatch: pytest.MonkeyPatch, + general_settings: dict[str, object], + requested_model: str | None, + creating_group: str, + expected: str, +) -> None: + from litellm.types.videos.utils import encode_video_id_with_provider + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + video_id: Final = encode_video_id_with_provider( + video_id="video_raw", + provider="azure", + model_id="sora-2", + deployment_id=f"dep-{creating_group}", + requested_model=requested_model, + ) + + assert ( + get_model_from_request( + request_data={"video_id": video_id}, + route="/v1/videos/{video_id}", + llm_router=_sora_fallback_router(), + ) + == expected + ) _BATCH_DEPLOYMENT_ID = "8d0eaa7e6c6f54a425dfd0062cb6b0dc" @@ -4463,9 +4671,7 @@ class TestIsRequestBodySafeBlocksAwsIdentitySelectors: ), ], ) -def test_configurable_clientside_parameters( - allowed_param, input_value, should_return_true -): +def test_configurable_clientside_parameters(allowed_param, input_value, should_return_true): router = Router( model_list=[ { @@ -4499,9 +4705,7 @@ def test_get_customer_user_header_from_mapping_returns_customer_header_with_mixe def test_get_customer_user_header_from_mapping_no_customer_returns_none(): from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping - mappings = [ - {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"} - ] + mappings = [{"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}] result = get_customer_user_header_from_mapping(mappings) assert result is None @@ -4529,17 +4733,13 @@ def test_get_internal_user_header_from_mapping_returns_internal_header(): def test_get_internal_user_header_from_mapping_no_internal_returns_none(): from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup - mappings = [ - {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"} - ] + mappings = [{"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}] result = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(mappings) assert result is None # Also support single mapping dict single_mapping = {"header_name": "X-Only-Customer", "litellm_user_role": "customer"} - result = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping( - single_mapping - ) + result = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(single_mapping) assert result is None @@ -4561,9 +4761,7 @@ def test_get_internal_user_header_from_mapping_no_internal_returns_none(): def test_get_model_from_request(request_data, expected_model): from litellm.proxy.auth.auth_utils import get_model_from_request - request_data = { - "target_model_names": "gpt-3.5-turbo, gpt-4o-mini-general-deployment" - } + request_data = {"target_model_names": "gpt-3.5-turbo, gpt-4o-mini-general-deployment"} route = "/openai/deployments/gpt-3.5-turbo" model = get_model_from_request(request_data, "/v1/files") assert model == ["gpt-3.5-turbo", "gpt-4o-mini-general-deployment"] @@ -4625,9 +4823,7 @@ def test_get_model_from_request(request_data, expected_model): ), ], ) -def test_get_model_from_request_vertex_ai_passthrough( - request_data, route, expected_model -): +def test_get_model_from_request_vertex_ai_passthrough(request_data, route, expected_model): """Test that get_model_from_request correctly extracts Vertex AI model from URL""" from litellm.proxy.auth.auth_utils import get_model_from_request diff --git a/tests/unit/proxy/video_endpoints/test_endpoints.py b/tests/unit/proxy/video_endpoints/test_endpoints.py index 5052d75f9f9..9809d26826b 100644 --- a/tests/unit/proxy/video_endpoints/test_endpoints.py +++ b/tests/unit/proxy/video_endpoints/test_endpoints.py @@ -28,7 +28,8 @@ mock's call args), and a brand-new kwarg added to this layer surfaces as a failu from contextlib import ExitStack from dataclasses import dataclass -from typing import Any, Dict, Optional +from collections.abc import Awaitable, Callable +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import orjson @@ -41,7 +42,9 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.utils import ProxyLogging from litellm.router import Router +from litellm.types.videos.main import VideoObject from litellm.types.videos.utils import ( + decode_video_id_with_provider, encode_character_id_with_provider, encode_video_id_with_provider, ) @@ -70,6 +73,29 @@ RESOLVED_MODELS: Dict[str, str] = {VIDEO_MODEL_ID: "azure-sora"} SENTINEL = object() +def _video_with_empty_model_id( + hidden_params: dict[str, object] | None = None, +) -> VideoObject: + video: Final[VideoObject] = VideoObject( + id=encode_video_id_with_provider("video_new456", "azure", ""), + object="video", + status="queued", + ) + video.hidden_params = dict(hidden_params or {}) + return video + + +def _assert_id_resolves_to_deployment(response: object, *, deployment_id: str | None = None) -> None: + assert isinstance(response, VideoObject) + expected_deployment: Final = {} if deployment_id is None else {"deployment_id": deployment_id} + assert decode_video_id_with_provider(response.id) == { + "custom_llm_provider": "azure", + "model_id": VIDEO_MODEL_ID, + "video_id": "video_new456", + **expected_deployment, + } + + class FakeRequest: """Minimal stand-in. headers/query_params are read by the provider readers (mocked) and on the edit path the raw body is parsed for real via orjson.""" @@ -105,6 +131,9 @@ class Harness: assert self.base_process.call_count == 1 return dict(self.base_process.call_args.args[0].data) + def respond_with(self, response: object) -> None: + self.base_process.return_value = response # pyright: ignore[reportAny] # MagicMock subclasses Any in typeshed + def route_type(self) -> str: return self.base_process.call_args.kwargs["route_type"] @@ -118,6 +147,7 @@ def harness(): side_effect=lambda model_id, custom_llm_provider=None: RESOLVED_MODELS.get(model_id) ) router.resolve_model_name_from_model_id = resolve_model + router.get_deployment.return_value = None read_body = AsyncMock(return_value={}) batch_to_bytesio = AsyncMock(return_value=[b"filebytes"]) @@ -244,6 +274,36 @@ async def test_generation__input_reference_attached(harness): } +@pytest.mark.asyncio +async def test_generation__restamps_id_with_asked_for_model_over_router_deployment_id(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id(hidden_params={"model_id": "deployment-uuid"})) + + resp: Final = await call_generation(harness, body={"model": VIDEO_MODEL_ID, "prompt": "x"}) + + _assert_id_resolves_to_deployment(resp, deployment_id="deployment-uuid") + + +@pytest.mark.asyncio +async def test_generation__restamps_id_with_request_model_when_no_hidden_model_id(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id()) + + resp: Final = await call_generation(harness, body={"model": VIDEO_MODEL_ID, "prompt": "x"}) + + _assert_id_resolves_to_deployment(resp) + + +@pytest.mark.asyncio +async def test_generation__unroutable_request_model_falls_back_to_router_deployment_id(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id(hidden_params={"model_id": "deployment-uuid"})) + + resp: Final = await call_generation(harness, body={"model": "team-public-sora", "prompt": "x"}) + + assert isinstance(resp, VideoObject) + decoded: Final = decode_video_id_with_provider(resp.id) + assert decoded.get("model_id") == "deployment-uuid" + assert decoded.get("deployment_id") == "deployment-uuid" + + @pytest.mark.asyncio async def test_generation__exception_routed_through_handler(harness): harness.base_process.side_effect = ValueError("provider boom") @@ -284,6 +344,24 @@ async def test_status__model_encoded_id_full_contract(harness): } +@pytest.mark.asyncio +async def test_status__restamps_response_id_with_decoded_model_id(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id()) + + resp: Final = await call_status(harness, AZURE_VIDEO_ID) + + _assert_id_resolves_to_deployment(resp) + + +@pytest.mark.asyncio +async def test_status__id_without_a_deployment_takes_the_deployment_that_answered_the_poll(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id(hidden_params={"model_id": "deployment-uuid"})) + + resp: Final = await call_status(harness, AZURE_VIDEO_ID) + + _assert_id_resolves_to_deployment(resp, deployment_id="deployment-uuid") + + @pytest.mark.asyncio async def test_status__plain_id_defaults_to_openai(harness): await call_status(harness, "video_plain") @@ -387,6 +465,30 @@ async def test_status_and_content__hand_the_served_member_to_the_request_process assert harness.base_process.call_args.args[0].served_model == "wan3-b" +@pytest.mark.asyncio +@pytest.mark.parametrize("call", [call_status, call_content]) +async def test_status_and_content__shared_provider_model_routes_to_the_creating_deployments_group( + harness: Harness, call: Callable[[Harness, str], Awaitable[object]] +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": "azure/sora-2", "api_base": f"https://{name}.example.com", "api_key": name}, + "model_info": {"id": f"dep-{name}"}, + } + for name in ("sora-team", "sora-shared") + ] + ) + harness.base_process.return_value = b"x" + video_id: Final = encode_video_id_with_provider("video_raw", "azure", "sora-2", deployment_id="dep-sora-shared") + + with patch.object(proxy_server, "llm_router", router): + await call(harness, video_id) + + assert harness.processor_data()["model"] == "sora-shared" + + # =========================================================================== # # POST /v1/videos/edits - video_edit # # =========================================================================== # @@ -422,6 +524,15 @@ async def test_edit__extracts_nested_video_id_full_contract(harness): } +@pytest.mark.asyncio +async def test_edit__restamps_new_video_id_with_source_model_id(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id()) + + resp: Final = await call_edit(harness, body={"prompt": "brighter", "video": {"id": AZURE_VIDEO_ID}}) + + _assert_id_resolves_to_deployment(resp) + + @pytest.mark.asyncio async def test_edit__provider_from_body_data_for_plain_id(harness): """For a plain id, get_custom_provider_from_data (run for real) pulls the @@ -473,6 +584,23 @@ async def test_edit__json_string_video_reference_from_form_field(harness): assert harness.processor_data()["video_id"] == "video_plain" +@pytest.mark.asyncio +async def test_edit__uploaded_video_restamps_new_id_with_requested_model(harness: Harness) -> None: + import io + + upload: Final = StarletteUploadFile(file=io.BytesIO(b"rawmp4"), filename="clip.mp4") + harness.read_body.return_value = {"model": VIDEO_MODEL_ID, "prompt": "make it nighttime", "video": upload} + harness.respond_with(_video_with_empty_model_id()) + + resp: Final = await endpoints.video_edit( + request=FakeRequest(raw_body=b"multipart"), + fastapi_response=Response(), + user_api_key_dict=_user(), + ) + + _assert_id_resolves_to_deployment(resp) + + @pytest.mark.asyncio async def test_edit__uploaded_video_file_is_forwarded_not_dropped(harness): """A multipart-uploaded source video must be converted to bytes and attached @@ -565,6 +693,15 @@ async def test_remix__model_encoded_id_full_contract(harness): } +@pytest.mark.asyncio +async def test_remix__restamps_new_video_id_with_source_model_id(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id()) + + resp: Final = await call_remix(harness, AZURE_VIDEO_ID, body={"prompt": "new colors"}) + + _assert_id_resolves_to_deployment(resp) + + @pytest.mark.asyncio async def test_remix__provider_from_body_data_not_request_body_reader(harness): """remix resolves the provider from data.get('custom_llm_provider'), never @@ -694,6 +831,20 @@ async def test_get_character__plain_id_defaults_openai_no_encode(harness): assert resp["id"] == "char_raw3" +@pytest.mark.asyncio +async def test_get_character__model_served_only_by_a_provider_wildcard_routes_through_it(harness: Harness) -> None: + router: Final = Router( + model_list=[{"model_name": "openai/*", "litellm_params": {"model": "openai/*", "api_key": "key-w"}}] + ) + harness.base_process.return_value = {"id": "char_raw4"} + character_id: Final = encode_character_id_with_provider("char_orig", "openai", "sora-2") + + with patch.object(proxy_server, "llm_router", router): + await call_get_character(harness, character_id) + + assert harness.processor_data()["model"] == "sora-2" + + # =========================================================================== # # POST /v1/videos/extensions - video_extension # # =========================================================================== # @@ -723,3 +874,12 @@ async def test_extension__extracts_nested_video_id_full_contract(harness): "custom_llm_provider": "azure", "model": "azure-sora", } + + +@pytest.mark.asyncio +async def test_extension__restamps_new_video_id_with_source_model_id(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id()) + + resp: Final = await call_extension(harness, body={"prompt": "continue", "video": {"id": AZURE_VIDEO_ID}}) + + _assert_id_resolves_to_deployment(resp) diff --git a/tests/unit/proxy/video_endpoints/test_utils.py b/tests/unit/proxy/video_endpoints/test_utils.py index 9a2c208c075..01a2b7c2840 100644 --- a/tests/unit/proxy/video_endpoints/test_utils.py +++ b/tests/unit/proxy/video_endpoints/test_utils.py @@ -1,11 +1,12 @@ """ Pure-logic contract tests for litellm/proxy/video_endpoints/utils.py -Four helpers the video proxy endpoints lean on: +Five helpers the video proxy endpoints lean on: - extract_model_from_target_model_names: first model from a comma string / list - video_reference_to_id: normalize a video reference (dict / bare id / JSON string) to an id - get_custom_provider_from_data: provider precedence (top-level > extra_body) - encode_character_id_in_response: re-encode a response id in place + - encode_video_id_in_response: re-stamp a VideoObject id with the resolved model_id Every test asserts the exact result (or identity), so a mutation that flips a branch, drops a strip/filter, or changes precedence fails. The only collaborator @@ -13,19 +14,29 @@ is encode_character_id_with_provider, which runs for real; encoding assertions are checked by the genuine decode round-trip. """ +from typing import Final import pytest +from litellm.router import Router +from litellm.types.videos.main import DecodedVideoId from litellm.proxy.video_endpoints.utils import ( encode_character_id_in_response, + encode_video_id_in_response, extract_model_from_target_model_names, + fallback_requested_model, get_custom_provider_from_data, + routable_model_name, + video_model_group, video_reference_to_id, ) +from litellm.types.videos.main import VideoObject from litellm.types.videos.utils import ( decode_character_id_with_provider, + decode_video_id_with_provider, encode_character_id_with_provider, + encode_video_id_with_provider, ) # =========================================================================== # @@ -103,12 +114,7 @@ def test_provider__falsy_top_level_falls_through_to_extra_body(falsy): def test_provider__from_extra_body_dict(): - assert ( - get_custom_provider_from_data( - {"extra_body": {"custom_llm_provider": "bedrock"}} - ) - == "bedrock" - ) + assert get_custom_provider_from_data({"extra_body": {"custom_llm_provider": "bedrock"}}) == "bedrock" def test_provider__from_extra_body_json_string(): @@ -126,10 +132,7 @@ def test_provider__json_string_parsing_to_non_dict_is_none(): def test_provider__extra_body_provider_not_a_string_is_none(): - assert ( - get_custom_provider_from_data({"extra_body": {"custom_llm_provider": 123}}) - is None - ) + assert get_custom_provider_from_data({"extra_body": {"custom_llm_provider": 123}}) is None @pytest.mark.parametrize( @@ -161,9 +164,7 @@ def test_encode__dict_with_id_mutates_in_place_and_preserves_other_keys(): assert out is response # same dict, mutated in place assert out["object"] == "character" and out["name"] == "hero" - assert out["id"] == encode_character_id_with_provider( - "char_raw", "azure", "model-1" - ) + assert out["id"] == encode_character_id_with_provider("char_raw", "azure", "model-1") decoded = decode_character_id_with_provider(out["id"]) assert decoded["custom_llm_provider"] == "azure" assert decoded["model_id"] == "model-1" @@ -206,3 +207,387 @@ def test_encode__object_without_id_attr_returned_unchanged(): out = encode_character_id_in_response(resp, "azure", "model-1") assert out is resp assert not hasattr(resp, "id") + + +# =========================================================================== # +# encode_video_id_in_response +# =========================================================================== # + + +def _video(video_id: str, hidden_params: dict[str, object] | None = None) -> VideoObject: + video: Final[VideoObject] = VideoObject(id=video_id, object="video", status="queued") + video.hidden_params = dict(hidden_params or {}) + return video + + +def test_encode_video__fills_empty_model_id_and_deployment_from_hidden_params() -> None: + original_id: Final = encode_video_id_with_provider("video_raw", "openai", "") + video: Final = _video(original_id, hidden_params={"model_id": "deployment-1"}) + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id=None) + + assert isinstance(out, VideoObject) + assert decode_video_id_with_provider(out.id) == { + "custom_llm_provider": "openai", + "model_id": "deployment-1", + "video_id": "video_raw", + "deployment_id": "deployment-1", + } + assert out is not video + assert video.id == original_id + assert out.hidden_params == {"model_id": "deployment-1"} + assert out.status == "queued" + + +def test_encode_video__asked_for_model_wins_over_hidden_deployment_id() -> None: + video: Final = _video( + encode_video_id_with_provider("video_raw", "openai", ""), hidden_params={"model_id": "deployment-1"} + ) + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id="sora-alias") + + assert isinstance(out, VideoObject) + decoded: Final = decode_video_id_with_provider(out.id) + assert decoded.get("model_id") == "sora-alias" + assert decoded.get("deployment_id") == "deployment-1" + + +def test_encode_video__stamped_id_keeps_model_and_served_group_and_gains_routed_deployment() -> None: + original_id: Final = encode_video_id_with_provider("task_1", "volcengine", "seedance-group", "seedance-east") + video: Final = _video(original_id, hidden_params={"model_id": "deployment-1"}) + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id="other") + + assert isinstance(out, VideoObject) + assert decode_video_id_with_provider(out.id) == { + "custom_llm_provider": "volcengine", + "model_id": "seedance-group", + "video_id": "task_1", + "served_model": "seedance-east", + "deployment_id": "deployment-1", + } + + +def test_encode_video__routed_deployment_replaces_the_one_in_the_id() -> None: + original_id: Final = encode_video_id_with_provider("video_raw", "openai", "sora", deployment_id="deployment-old") + video: Final = _video(original_id, hidden_params={"model_id": "deployment-new"}) + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id=None) + + assert isinstance(out, VideoObject) + assert decode_video_id_with_provider(out.id).get("deployment_id") == "deployment-new" + + +@pytest.mark.parametrize( + ("polled_deployment", "answering_deployment", "returned_deployment"), + [ + (None, "deployment-routed", "deployment-routed"), + ("deployment-gone", "deployment-routed", "deployment-routed"), + ("deployment-creator", None, "deployment-creator"), + ], +) +def test_encode_video__a_poll_takes_the_deployment_that_answered_else_keeps_the_one_its_id_named( + polled_deployment: str | None, + answering_deployment: str | None, + returned_deployment: str, +) -> None: + polled_id: Final = encode_video_id_with_provider("video_raw", "openai", "sora", deployment_id=polled_deployment) + video: Final = _video( + encode_video_id_with_provider("video_raw", "openai", ""), + hidden_params={} if answering_deployment is None else {"model_id": answering_deployment}, + ) + + out: Final = encode_video_id_in_response( + video, + fallback_provider=None, + fallback_model_id="sora", + polled_video=decode_video_id_with_provider(polled_id), + ) + + assert isinstance(out, VideoObject) + assert out.id == encode_video_id_with_provider("video_raw", "openai", "sora", deployment_id=returned_deployment) + + +def test_encode_video__a_poll_keeps_the_group_a_fallback_create_was_serving() -> None: + polled_id: Final = encode_video_id_with_provider( + "video_raw", "azure", "sora-2", deployment_id="dep-backup", requested_model="sora-pair" + ) + video: Final = _video(encode_video_id_with_provider("video_raw", "azure", "sora-2")) + + out: Final = encode_video_id_in_response( + video, fallback_provider=None, fallback_model_id=None, polled_video=decode_video_id_with_provider(polled_id) + ) + + assert isinstance(out, VideoObject) + assert out.id == polled_id + + +def test_encode_video__a_remix_of_a_fallback_video_keeps_its_requested_group() -> None: + video: Final = _video( + encode_video_id_with_provider("video_new", "azure", "sora-2"), hidden_params={"model_id": "dep-backup"} + ) + + out: Final = encode_video_id_in_response( + video, fallback_provider=None, fallback_model_id=None, requested_model="sora-pair" + ) + + assert isinstance(out, VideoObject) + assert decode_video_id_with_provider(out.id).get("requested_model") == "sora-pair" + + +def test_encode_video__id_already_naming_the_routed_deployment_is_unchanged() -> None: + original_id: Final = encode_video_id_with_provider("video_raw", "openai", "sora", deployment_id="deployment-1") + video: Final = _video(original_id, hidden_params={"model_id": "deployment-1"}) + + assert encode_video_id_in_response(video, fallback_provider=None, fallback_model_id="other") is video + + +def test_encode_video__unrouted_response_keeps_the_deployment_in_its_id() -> None: + original_id: Final = encode_video_id_with_provider("video_raw", "openai", "", deployment_id="deployment-1") + + out: Final = encode_video_id_in_response(_video(original_id), fallback_provider=None, fallback_model_id="sora") + + assert isinstance(out, VideoObject) + decoded: Final = decode_video_id_with_provider(out.id) + assert decoded.get("model_id") == "sora" + assert decoded.get("deployment_id") == "deployment-1" + + +def test_encode_video__name_with_id_separator_is_never_written_into_the_id() -> None: + video: Final = _video( + encode_video_id_with_provider("video_raw", "openai", ""), hidden_params={"model_id": "deployment-1"} + ) + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id="sora;model_id:evil") + + assert isinstance(out, VideoObject) + assert decode_video_id_with_provider(out.id) == { + "custom_llm_provider": "openai", + "model_id": "deployment-1", + "video_id": "video_raw", + "deployment_id": "deployment-1", + } + + +def test_encode_video__falls_back_to_given_model_id() -> None: + video: Final = _video(encode_video_id_with_provider("video_raw", "openai", "")) + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id="sora-2") + + assert isinstance(out, VideoObject) + decoded: Final = decode_video_id_with_provider(out.id) + assert decoded.get("model_id") == "sora-2" + assert decoded.get("video_id") == "video_raw" + + +def test_encode_video__keeps_existing_model_id_over_fallback() -> None: + original_id: Final = encode_video_id_with_provider("video_raw", "azure", "deployment-1") + + out: Final = encode_video_id_in_response(_video(original_id), fallback_provider="openai", fallback_model_id="other") + + assert isinstance(out, VideoObject) + assert out.id == original_id + + +def test_encode_video__plain_id_encoded_with_fallback_provider() -> None: + video: Final = _video("video_raw", hidden_params={"model_id": "deployment-1"}) + + out: Final = encode_video_id_in_response(video, fallback_provider="openai", fallback_model_id=None) + + assert isinstance(out, VideoObject) + assert decode_video_id_with_provider(out.id) == { + "custom_llm_provider": "openai", + "model_id": "deployment-1", + "video_id": "video_raw", + "deployment_id": "deployment-1", + } + + +def test_encode_video__no_model_id_anywhere_unchanged() -> None: + video: Final = _video(encode_video_id_with_provider("video_raw", "openai", "")) + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id=None) + + assert out is video + + +def test_encode_video__no_provider_anywhere_unchanged() -> None: + video: Final = _video("video_raw") + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id="sora-2") + + assert out is video + + +@pytest.mark.parametrize("response", [b"video-bytes", {"id": "video_raw"}, None]) +def test_encode_video__non_video_object_passthrough(response: object) -> None: + assert encode_video_id_in_response(response, "openai", "sora-2") is response + + +def _router_with_groups(*names: str) -> Router: + return Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/sora-2", "api_key": "group-key"}} for name in names + ] + ) + + +def test_routable_model_name__keeps_a_name_the_router_resolves() -> None: + assert routable_model_name("sora-alias", _router_with_groups("sora-alias"), "openai") == "sora-alias" + + +def test_routable_model_name__drops_a_name_the_router_cannot_route() -> None: + assert routable_model_name("team-public-name", _router_with_groups("sora-alias"), "openai") is None + + +def test_routable_model_name__keeps_the_name_without_a_router() -> None: + assert routable_model_name("sora-2", None, None) == "sora-2" + + +@pytest.mark.parametrize("model", [None, "", 7, "sora;model_id:evil"]) +def test_routable_model_name__never_returns_a_name_unsafe_for_the_id(model: object) -> None: + assert routable_model_name(model, _router_with_groups("sora;model_id:evil"), "openai") is None + + +def _two_groups_sharing_sora() -> Router: + return Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": "azure/sora-2", "api_base": f"https://{name}.example.com", "api_key": name}, + "model_info": {"id": f"dep-{name}"}, + } + for name in ("sora-team", "sora-shared") + ] + ) + + +def _decoded(model_id: str, deployment_id: str | None) -> DecodedVideoId: + base: Final = DecodedVideoId(custom_llm_provider="azure", model_id=model_id, video_id="video_raw") + return base if deployment_id is None else {**base, "deployment_id": deployment_id} + + +@pytest.mark.parametrize("deployment_id", ["dep-sora-team", "dep-sora-shared"]) +def test_video_model_group__shared_provider_model_resolves_to_the_creating_deployments_group( + deployment_id: str, +) -> None: + router: Final = _two_groups_sharing_sora() + + assert video_model_group(_decoded("sora-2", deployment_id), router) == deployment_id.removeprefix("dep-") + + +@pytest.mark.parametrize("deployment_id", [None, "dep-gone"]) +def test_video_model_group__without_a_known_deployment_matches_plain_resolution(deployment_id: str | None) -> None: + router: Final = _two_groups_sharing_sora() + + assert video_model_group(_decoded("sora-2", deployment_id), router) == router.resolve_model_name_from_model_id( + "sora-2", "azure" + ) + + +def test_video_model_group__keeps_a_group_named_in_the_id_over_a_deployment_of_another_group() -> None: + assert video_model_group(_decoded("sora-team", "dep-sora-shared"), _two_groups_sharing_sora()) == "sora-team" + + +@pytest.mark.parametrize("deployment_id", ["dep-sora-team", "dep-sora-shared"]) +def test_video_model_group__keeps_an_alias_of_the_creating_deployments_group(deployment_id: str) -> None: + router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": f"openai/sora-{name}", "api_key": name}, + "model_info": {"id": f"dep-{name}"}, + } + for name in ("sora-team", "sora-shared") + ], + model_group_alias={"sora-alias": deployment_id.removeprefix("dep-")}, + ) + provider_model: Final = f"sora-{deployment_id.removeprefix('dep-')}" + + assert video_model_group(_decoded(provider_model, deployment_id), router) == "sora-alias" + + +def test_video_model_group__names_a_team_model_by_its_public_name() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "model_name_team-1_internal", + "litellm_params": {"model": "azure/sora-2", "api_key": "team-key"}, + "model_info": {"id": "dep-team", "team_id": "team-1", "team_public_model_name": "team-sora"}, + } + ] + ) + + assert video_model_group(_decoded("sora-2", "dep-team"), router) == "team-sora" + + +def _fallback_router() -> Router: + return Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": "azure/sora-2", "api_base": f"https://{name}.example.com", "api_key": name}, + "model_info": {"id": f"dep-{name}"}, + } + for name in ("sora-pair", "sora-backup", "sora-last") + ], + model_group_alias={"sora-alias": "sora-pair"}, + fallbacks=[{"sora-pair": ["sora-backup"]}, {"sora-backup": ["sora-last"]}], + ) + + +def _wildcard_router() -> Router: + return Router( + model_list=[ + { + "model_name": "gemini/*", + "litellm_params": {"model": "gemini/*", "api_key": "wildcard-key"}, + "model_info": {"id": "dep-wildcard"}, + } + ] + ) + + +@pytest.mark.parametrize( + ("requested_model", "answering_deployment", "expected"), + [ + ("sora-pair", "dep-sora-backup", "sora-pair"), + ("sora-pair", "dep-sora-last", "sora-pair"), + ("sora-backup", "dep-sora-pair", None), + ("sora-pair", "dep-sora-pair", None), + ("sora-alias", "dep-sora-pair", None), + ("sora-pair", None, None), + (None, "dep-sora-backup", None), + ], + ids=[ + "fallback-served", + "chained-fallback-served", + "served-by-a-group-it-never-falls-back-to", + "served-by-own-group", + "served-by-alias-target", + "unrouted", + "no-requested-model", + ], +) +def test_fallback_requested_model__names_the_asked_group_only_when_another_group_served_the_create( + requested_model: str | None, answering_deployment: str | None, expected: str | None +) -> None: + video: Final = _video( + encode_video_id_with_provider("video_raw", "azure", "sora-2"), + hidden_params={} if answering_deployment is None else {"model_id": answering_deployment}, + ) + + assert fallback_requested_model(requested_model, video, _fallback_router()) == expected + + +def test_fallback_requested_model__is_none_for_a_create_a_provider_wildcard_served() -> None: + video: Final = _video( + encode_video_id_with_provider("operations/op_1", "gemini", "veo-3.1-fast-generate-preview"), + hidden_params={"model_id": "dep-wildcard"}, + ) + + assert fallback_requested_model("gemini/veo-3.1-fast-generate-preview", video, _wildcard_router()) is None + + +def test_video_model_group__is_none_without_a_router() -> None: + assert video_model_group(_decoded("sora-2", "dep-sora-shared"), None) is None diff --git a/tests/unit/test_router/test_router_provider_endpoints.py b/tests/unit/test_router/test_router_provider_endpoints.py index 708322f4fec..9312204f471 100644 --- a/tests/unit/test_router/test_router_provider_endpoints.py +++ b/tests/unit/test_router/test_router_provider_endpoints.py @@ -16,6 +16,7 @@ from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.proxy._types import UserAPIKeyAuth from litellm.router_utils.router_callbacks.track_deployment_metrics import ( get_deployment_successes_for_current_minute, ) @@ -23,6 +24,7 @@ from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import HttpxBinaryResponseContent from litellm.types.router import RoutingStrategy from litellm.types.utils import ImageResponse, ModelResponse, RerankResponse +from litellm.types.videos.utils import encode_video_id_with_provider from openai import AsyncAzureOpenAI, AsyncOpenAI EVENT_TIMEOUT_SECONDS: Final = 5 @@ -730,3 +732,248 @@ async def test_router_moderation_endpoint_factory_invokes_default_model( assert route.call_count == 1 assert json.loads(route.calls[0].request.content) == {"input": "hello"} assert response.id == "modr-2" + + +def _sora_deployment( + model_name: str, api_key: str, deployment_id: str, access_groups: tuple[str, ...] = () +) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": {"model": "openai/sora-2", "api_key": api_key}, + "model_info": {"id": deployment_id, "access_groups": list(access_groups)}, + } + + +def _two_account_sora_router() -> Router: + return Router( + model_list=[ + _sora_deployment("sora", "key-a", "dep-a"), + _sora_deployment("sora", "key-b", "dep-b"), + _sora_deployment("other", "key-x", "dep-x"), + ], + model_group_alias={"sora-alias": "sora"}, + ) + + +async def _poll_status_keys( + router: Router, + model: str, + deployment_id: str, + respx_mock: respx.MockRouter, + caller: UserAPIKeyAuth | None = None, +) -> set[str]: + route: Final = respx_mock.get("https://api.openai.com/v1/videos/video_raw").respond( + 200, json={"id": "video_raw", "object": "video", "status": "in_progress"} + ) + video_id: Final = encode_video_id_with_provider("video_raw", "openai", "sora", deployment_id=deployment_id) + metadata: Final = {} if caller is None else {"user_api_key_auth": caller} + for _ in range(10): + await router.avideo_status(model=model, video_id=video_id, litellm_metadata=dict(metadata)) + return {call.request.headers["authorization"] for call in route.calls} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["sora", "sora-alias"]) +async def test_video_status_polls_the_deployment_that_created_the_video( + model: str, respx_mock: respx.MockRouter +) -> None: + keys: Final = await _poll_status_keys(_two_account_sora_router(), model, "dep-b", respx_mock) + + assert keys == {"Bearer key-b"} + + +def test_sync_video_status_polls_the_deployment_that_created_the_video(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.get("https://api.openai.com/v1/videos/video_raw").respond( + 200, json={"id": "video_raw", "object": "video", "status": "in_progress"} + ) + router: Final = _two_account_sora_router() + video_id: Final = encode_video_id_with_provider("video_raw", "openai", "sora", deployment_id="dep-b") + + for _ in range(10): + router.video_status(model="sora", video_id=video_id) + + assert {call.request.headers["authorization"] for call in route.calls} == {"Bearer key-b"} + + +@pytest.mark.asyncio +async def test_video_status_on_a_wildcard_polls_the_deployment_that_created_the_video( + respx_mock: respx.MockRouter, +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": key}, + "model_info": {"id": dep}, + } + for dep, key in (("wc-a", "key-a"), ("wc-b", "key-b")) + ] + ) + + keys: Final = await _poll_status_keys(router, "openai/sora-2", "wc-b", respx_mock) + + assert keys == {"Bearer key-b"} + + +@pytest.mark.asyncio +async def test_video_status_never_pins_to_a_deployment_in_cooldown(respx_mock: respx.MockRouter) -> None: + router: Final = _two_account_sora_router() + router.cooldown_cache.add_deployment_to_cooldown( + model_id="dep-b", + original_exception=litellm.RateLimitError(message="upstream 429", llm_provider="openai", model="sora-2"), + exception_status=429, + cooldown_time=60, + ) + + keys: Final = await _poll_status_keys(router, "sora", "dep-b", respx_mock) + + assert keys == {"Bearer key-a"} + + +@pytest.mark.asyncio +async def test_video_status_on_a_routing_group_polls_the_member_deployment_that_created_the_video( + respx_mock: respx.MockRouter, +) -> None: + router: Final = Router( + model_list=[_sora_deployment("sora-a", "key-a", "dep-a"), _sora_deployment("sora-b", "key-b", "dep-b")], + routing_groups=[ + {"group_name": "sora-group", "models": ["sora-a", "sora-b"], "routing_strategy": "simple-shuffle"} + ], + ) + + keys: Final = await _poll_status_keys(router, "sora-group", "dep-b", respx_mock) + + assert keys == {"Bearer key-b"} + + +@pytest.mark.asyncio +async def test_video_status_never_pins_to_a_deployment_outside_the_requested_group( + respx_mock: respx.MockRouter, +) -> None: + keys: Final = await _poll_status_keys(_two_account_sora_router(), "sora", "dep-x", respx_mock) + + assert keys <= {"Bearer key-a", "Bearer key-b"} + + +@pytest.mark.asyncio +async def test_video_status_with_a_removed_deployment_routes_across_the_group(respx_mock: respx.MockRouter) -> None: + keys: Final = await _poll_status_keys(_two_account_sora_router(), "sora", "dep-gone", respx_mock) + + assert keys <= {"Bearer key-a", "Bearer key-b"} + assert keys + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("key_access_group", "expected_key"), [("ag-a", "Bearer key-a"), ("ag-b", "Bearer key-b")]) +async def test_video_status_pins_only_to_a_deployment_in_the_callers_access_groups( + key_access_group: str, expected_key: str, respx_mock: respx.MockRouter +) -> None: + router: Final = Router( + model_list=[ + _sora_deployment("sora", "key-a", "dep-a", access_groups=("ag-a",)), + _sora_deployment("sora", "key-b", "dep-b", access_groups=("ag-b",)), + ] + ) + + keys: Final = await _poll_status_keys( + router, "sora", "dep-b", respx_mock, caller=UserAPIKeyAuth(models=[key_access_group]) + ) + + assert keys == {expected_key} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["sora", "sora-alias"]) +async def test_pinned_video_status_logs_the_requested_model_group( + model: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + capture: Final = _RouterLoggingCapture("dep-b") + monkeypatch.setattr(litellm, "callbacks", [capture]) + + keys: Final = await _poll_status_keys(_two_account_sora_router(), model, "dep-b", respx_mock) + _, standard_logging_object = await capture.next_event() + + assert keys == {"Bearer key-b"} + assert isinstance(standard_logging_object, dict) + assert standard_logging_object["model_group"] == model + + +def _ordered_chat_router() -> Router: + return Router( + model_list=[ + { + "model_name": "chat-ordered", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": key, "order": order}, + "model_info": {"id": deployment_id}, + } + for deployment_id, key, order in (("dep-primary", "key-a", 1), ("dep-backup", "key-b", 2)) + ] + ) + + +@pytest.mark.asyncio +async def test_a_video_id_in_a_chat_body_never_overrides_deployment_order() -> None: + router: Final = _ordered_chat_router() + video_id: Final = encode_video_id_with_provider("video_raw", "openai", "chat-ordered", deployment_id="dep-backup") + + chats: Final = [ + await router.acompletion( + model="chat-ordered", messages=list(ROUTING_MESSAGES), mock_response="ok", video_id=video_id + ) + for _ in range(4) + ] + responses: Final = [ + await router.aresponses(model="chat-ordered", input="route this request", mock_response="ok", video_id=video_id) + for _ in range(4) + ] + + assert {chat._hidden_params["model_id"] for chat in chats} == {"dep-primary"} + assert {response._hidden_params["model_id"] for response in responses} == {"dep-primary"} + + +@pytest.mark.asyncio +async def test_video_lookup_404_never_cools_down_the_deployment_that_answered(respx_mock: respx.MockRouter) -> None: + respx_mock.get("https://api.openai.com/v1/videos/video_gone").respond( + 404, json={"error": {"code": "NotFound", "message": "Video not found"}} + ) + router: Final = Router( + model_list=[_sora_deployment("sora", "key-a", "dep-a"), _sora_deployment("sora", "key-b", "dep-b")], + num_retries=0, + allowed_fails=1, + ) + video_id: Final = encode_video_id_with_provider("video_gone", "openai", "sora", deployment_id="dep-b") + + for _ in range(3): + with pytest.raises(litellm.NotFoundError): + await router.avideo_status(model="sora", video_id=video_id) + + assert await router.cooldown_cache.async_get_active_cooldowns(["dep-a", "dep-b"], parent_otel_span=None) == [] + + +@pytest.mark.asyncio +async def test_chat_404_still_cools_down_the_deployment(respx_mock: respx.MockRouter) -> None: + respx_mock.post("https://api.openai.com/v1/chat/completions").respond( + 404, json={"error": {"code": "NotFound", "message": "Deployment not found"}} + ) + router: Final = Router( + model_list=[ + { + "model_name": "chat", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "key-a"}, + "model_info": {"id": "dep-chat"}, + }, + { + "model_name": "chat", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "key-b"}, + "model_info": {"id": "dep-chat-b"}, + }, + ], + num_retries=0, + allowed_fails=1, + ) + + for _ in range(3): + with contextlib.suppress(litellm.NotFoundError, litellm.RateLimitError): + await router.acompletion(model="chat", messages=list(ROUTING_MESSAGES)) + + assert await router.cooldown_cache.async_get_active_cooldowns(["dep-chat", "dep-chat-b"], parent_otel_span=None) diff --git a/tests/unit/test_video_generation.py b/tests/unit/test_video_generation.py index 5c4bc8ea652..df7e9c68b90 100644 --- a/tests/unit/test_video_generation.py +++ b/tests/unit/test_video_generation.py @@ -2,6 +2,7 @@ import asyncio import io import json import os +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1269,6 +1270,35 @@ def test_encode_video_id_with_provider_handles_azure_video_prefix(): assert encoded_twice == encoded_id # Should return the same encoded ID +@pytest.mark.parametrize( + "served_model, deployment_id, expected_extras", + [ + ("seedance-east", "dep-1", {"served_model": "seedance-east", "deployment_id": "dep-1"}), + (None, "dep-1", {"deployment_id": "dep-1"}), + ("seedance-east", None, {"served_model": "seedance-east"}), + (None, None, {}), + ], +) +def test_video_id_round_trips_served_model_and_deployment_id( + served_model: str | None, deployment_id: str | None, expected_extras: dict[str, str] +) -> None: + from litellm.types.videos.utils import ( + decode_video_id_with_provider, + encode_video_id_with_provider, + ) + + encoded: Final = encode_video_id_with_provider( + "task_1", "volcengine", "seedance-group", served_model=served_model, deployment_id=deployment_id + ) + + assert decode_video_id_with_provider(encoded) == { + "custom_llm_provider": "volcengine", + "model_id": "seedance-group", + "video_id": "task_1", + **expected_extras, + } + + class TestVideoListTransformation: """Tests for video list request/response transformation with provider ID encoding."""