diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 5cbc69669d7..dc245d42862 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -762,12 +762,19 @@ def _map_openai_like_exception( llm_provider=custom_llm_provider, model=model, ) - elif original_exception.status_code == 401 or original_exception.status_code == 403: + elif original_exception.status_code == 401: raise AuthenticationError( message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", llm_provider=custom_llm_provider, model=model, ) + elif original_exception.status_code == 403: + raise PermissionDeniedError( + message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", + llm_provider=custom_llm_provider, + model=model, + response=_response_or_stub(original_exception, status_code=403), + ) elif original_exception.status_code == 400: raise BadRequestError( message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", @@ -2194,6 +2201,120 @@ def _map_openrouter_exception( ) +def _response_or_stub(original_exception: _ProviderHTTPException, status_code: int) -> httpx.Response: + response: Final = original_exception.response if hasattr(original_exception, "response") else None + if response is not None: + return response + return httpx.Response( + status_code=status_code, request=httpx.Request(method="POST", url="https://docs.litellm.ai/docs") + ) + + +def _map_exception_by_status( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_provider: str, + extra_information: str, +) -> None: + status_code: Final = original_exception.status_code if hasattr(original_exception, "status_code") else None + if not isinstance(status_code, int) or status_code < 400: + return + message: Final = f"{exception_provider} - {error_str}" + response: Final = original_exception.response if hasattr(original_exception, "response") else None + match status_code: + case 401: + raise AuthenticationError( + message=message, + llm_provider=custom_llm_provider, + model=model, + response=response, + litellm_debug_info=extra_information, + ) + case 403: + raise PermissionDeniedError( + message=message, + llm_provider=custom_llm_provider, + model=model, + response=_response_or_stub(original_exception, status_code=status_code), + litellm_debug_info=extra_information, + ) + case 404: + raise NotFoundError( + message=message, + model=model, + llm_provider=custom_llm_provider, + response=response, + litellm_debug_info=extra_information, + ) + case 408: + raise Timeout( + message=message, + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) + case 429: + raise RateLimitError( + message=message, + model=model, + llm_provider=custom_llm_provider, + response=response, + litellm_debug_info=extra_information, + ) + case 500: + raise InternalServerError( + message=message, + llm_provider=custom_llm_provider, + model=model, + response=response, + litellm_debug_info=extra_information, + ) + case 502: + raise BadGatewayError( + message=message, + llm_provider=custom_llm_provider, + model=model, + response=response, + litellm_debug_info=extra_information, + ) + case 503: + raise ServiceUnavailableError( + message=message, + llm_provider=custom_llm_provider, + model=model, + response=response, + litellm_debug_info=extra_information, + ) + case 504: + raise Timeout( + message=message, + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + exception_status_code=status_code, + ) + case _ if status_code < 500: + raise BadRequestError( + message=message, + model=model, + llm_provider=custom_llm_provider, + response=response, + litellm_debug_info=extra_information, + ) + case _: + raise APIError( + status_code=status_code, + message=message, + llm_provider=custom_llm_provider, + model=model, + request=original_exception.request if hasattr(original_exception, "request") else None, + litellm_debug_info=extra_information, + ) + + def exception_type( model, original_exception, @@ -2508,6 +2629,14 @@ def exception_type( For unmapped exceptions - raise the exception with traceback - https://github.com/BerriAI/litellm/issues/4201 """ exception_mapping_worked = True + _map_exception_by_status( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_provider=exception_provider, + extra_information=extra_information, + ) if hasattr(original_exception, "request"): raise APIConnectionError( message=f"{exception_provider} - {error_str}", diff --git a/litellm/llms/together_ai/chat/transformation.py b/litellm/llms/together_ai/chat/transformation.py index 6bf53ee348b..d06bcb22119 100644 --- a/litellm/llms/together_ai/chat/transformation.py +++ b/litellm/llms/together_ai/chat/transformation.py @@ -4,8 +4,7 @@ Translates from OpenAI's `/v1/chat/completions` to Together AI's `/v1/chat/compl Docs: https://docs.together.ai/docs/chat-overview """ -from collections.abc import Container, Coroutine -from types import MappingProxyType +from collections.abc import Callable, Container, Coroutine from typing import ( Final, Literal, @@ -17,26 +16,42 @@ import litellm from litellm._logging import verbose_logger from litellm.exceptions import UnsupportedParamsError from litellm.types.llms.openai import AllMessageValues -from litellm.utils import supports_function_calling +from litellm.utils import supports_function_calling, supports_response_schema from ...openai.chat.gpt_transformation import OpenAIGPTConfig TOOL_CALLING_PARAMS: Final = ("tools", "tool_choice", "function_call") LITELLM_INTERNAL_ASSISTANT_FIELDS: Final = frozenset({"thinking_blocks", "provider_specific_fields"}) -PLAIN_TEXT_RESPONSE_FORMAT: Final = MappingProxyType({"type": "text"}) FUNCTION_CALLING_DOCS_URL: Final = "https://docs.together.ai/docs/function-calling" +STRUCTURED_OUTPUTS_DOCS_URL: Final = "https://docs.together.ai/docs/inference/chat/structured-outputs" + + +def _registry_verdict(model: str, flag: str, check: Callable[[str], bool]) -> bool | None: + try: + if check(model): + return True + except Exception as e: + verbose_logger.debug("Error checking together_ai %s for %s: %s", flag, model, e) + registry_entry: Final = litellm.model_cost.get(f"together_ai/{model}") + if isinstance(registry_entry, dict) and registry_entry.get(flag) is False: + return False + return None def _function_calling_verdict(model: str) -> bool | None: - try: - if supports_function_calling(model, custom_llm_provider="together_ai"): - return True - except Exception as e: - verbose_logger.debug("Error checking together_ai function calling support for %s: %s", model, e) - registry_entry: Final = litellm.model_cost.get(f"together_ai/{model}") - if isinstance(registry_entry, dict) and registry_entry.get("supports_function_calling") is False: - return False - return None + return _registry_verdict( + model, + "supports_function_calling", + lambda checked_model: supports_function_calling(checked_model, custom_llm_provider="together_ai"), + ) + + +def _response_schema_verdict(model: str) -> bool | None: + return _registry_verdict( + model, + "supports_response_schema", + lambda checked_model: supports_response_schema(checked_model, custom_llm_provider="together_ai"), + ) def _tool_params_to_drop(passed_params: Container[str], model: str, drop_params: bool) -> tuple[str, ...]: @@ -68,6 +83,32 @@ def _tool_params_to_drop(passed_params: Container[str], model: str, drop_params: ) +def _drop_response_format(passed_params: Container[str], model: str, drop_params: bool) -> bool: + if "response_format" not in passed_params: + return False + verdict: Final = _response_schema_verdict(model) + if verdict is True: + return False + if verdict is None: + verbose_logger.warning( + "together_ai model %s has no structured outputs entry in the model registry; passing response_format through for Together to validate. Docs - %s", + model, + STRUCTURED_OUTPUTS_DOCS_URL, + ) + return False + if drop_params or litellm.drop_params: + verbose_logger.warning( + "together_ai model %s does not support structured outputs per the model registry; dropping response_format. Docs - %s", + model, + STRUCTURED_OUTPUTS_DOCS_URL, + ) + return True + raise UnsupportedParamsError( + status_code=500, + message=f"together_ai does not support parameters: response_format, for model={model}. To drop it from the call, set `litellm.drop_params = True`.", + ) + + def _without_litellm_internal_fields(message: AllMessageValues) -> AllMessageValues: if message["role"] != "assistant" or LITELLM_INTERNAL_ASSISTANT_FIELDS.isdisjoint(message): return message @@ -112,18 +153,6 @@ class TogetherAIChatConfig(OpenAIGPTConfig): return super()._transform_messages(stripped, model, is_async=True) return super()._transform_messages(stripped, model, is_async=False) - def get_supported_openai_params(self, model: str) -> list: - supports_fc: Final = _function_calling_verdict(model) - supported_params: Final = super().get_supported_openai_params(model) - if supports_fc is True: - return supported_params - verbose_logger.debug( - "Only some together models support response_format. Docs - https://docs.together.ai/docs/function-calling" - ) - return [ # mutable-ok: the inherited contract returns a plain list; building fresh avoids mutating the base class's value - param for param in supported_params if param != "response_format" - ] - def map_openai_params( self, non_default_params: dict, @@ -134,6 +163,6 @@ class TogetherAIChatConfig(OpenAIGPTConfig): mapped_openai_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params) for param in _tool_params_to_drop(mapped_openai_params, model, drop_params): mapped_openai_params.pop(param) - if mapped_openai_params.get("response_format") == PLAIN_TEXT_RESPONSE_FORMAT: + if _drop_response_format(mapped_openai_params, model, drop_params): mapped_openai_params.pop("response_format") return mapped_openai_params diff --git a/litellm/proxy/common_utils/sse_keepalive.py b/litellm/proxy/common_utils/sse_keepalive.py index 429fff99ae6..26fccf8ee82 100644 --- a/litellm/proxy/common_utils/sse_keepalive.py +++ b/litellm/proxy/common_utils/sse_keepalive.py @@ -91,6 +91,17 @@ def is_sse_content_type(content_type: str | None) -> bool: return content_type is not None and content_type.split(";", 1)[0].strip().lower() == _SSE_MEDIA_TYPE +def split_complete_sse_frames(pending: bytes) -> tuple[bytes, bytes]: + """Split buffered SSE bytes into ``(complete_frames, unterminated_tail)``.""" + boundary_end: Final = max( + (pending.rfind(delimiter) + len(delimiter) for delimiter in _SSE_FRAME_DELIMITERS if delimiter in pending), + default=0, + ) + if boundary_end == 0: + return b"", pending + return pending[:boundary_end], pending[boundary_end:] + + def wrap_passthrough_sse_bytes_with_keepalive_pings( stream: AsyncGenerator[bytes, None], ping_interval_seconds: float | str | None, diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 3dafcc08551..21d12c8f720 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -28,6 +28,21 @@ if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth +_RESPONSES_API_PROVIDER_PREFIX: Final = "/openai" +_RESPONSES_API_CREATE_ROUTES: Final = frozenset({"/v1/responses", "/responses"}) + + +def _is_responses_api_create_route(request_route: str | None) -> bool: + if request_route is None: + return False + canonical: Final = ( + request_route[len(_RESPONSES_API_PROVIDER_PREFIX) :] + if request_route.startswith(_RESPONSES_API_PROVIDER_PREFIX + "/") + else request_route + ) + return canonical in _RESPONSES_API_CREATE_ROUTES + + class ResponsesIDSecurity(CustomLogger): def __init__(self): pass @@ -267,8 +282,7 @@ class ResponsesIDSecurity(CustomLogger): async for chunk in response: if ( isinstance(chunk, BaseLiteLLMOpenAIResponseObject) - and user_api_key_dict.request_route - == "/v1/responses" # only encrypt the response id for the responses api + and _is_responses_api_create_route(user_api_key_dict.request_route) and not general_settings.get("disable_responses_id_security", False) ): chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache) diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 4de6ef04d76..23cfef6576c 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -32,7 +32,7 @@ from __future__ import annotations import json import re -from collections.abc import Callable, Mapping, Sequence +from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from typing import ( TYPE_CHECKING, Final, @@ -43,7 +43,7 @@ from typing import ( from urllib.parse import quote, unquote from fastapi import HTTPException -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.managed_resources.isolation import ( @@ -52,6 +52,7 @@ from litellm.llms.base_llm.managed_resources.isolation import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit +from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames from litellm.repositories.table_repositories import ( ManagedFileRepository, ManagedObjectRepository, @@ -820,6 +821,121 @@ async def rewrite_response_ids( return mutated if changed else body +_RESPONSE_ID_PREFIX: Final = "resp_" +_STREAMED_RESPONSE_ID_SPEC: Final[_FieldSpec] = ("id", _RESPONSE_ID_PREFIX) +_SSE_DATA_PREFIX: Final = "data:" +_SSE_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, JsonValue]) + + +def _first_streamed_response(frames: bytes) -> tuple[str, Mapping[str, JsonValue]] | None: + for line in frames.decode("utf-8", errors="replace").splitlines(): + if not line.startswith(_SSE_DATA_PREFIX): + continue + try: + event = _SSE_EVENT_ADAPTER.validate_json(line[len(_SSE_DATA_PREFIX) :]) + except ValidationError: + continue + response = event.get("response") + if not isinstance(response, dict): + continue + raw_id = response.get("id") + if isinstance(raw_id, str) and raw_id.startswith(_RESPONSE_ID_PREFIX): + return raw_id, response + return None + + +class _StreamedResponseIdRewriter: + __slots__ = ("_is_create_route", "_pending", "_prisma_client", "_provider", "_replacement", "_user_api_key_dict") + + def __init__( + self, + provider: str, + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient, + is_create_route: bool, + ) -> None: + self._provider: Final = provider + self._user_api_key_dict: Final = user_api_key_dict + self._prisma_client: Final = prisma_client + self._is_create_route: Final = is_create_route + self._pending = b"" + self._replacement: tuple[bytes, bytes] | None = None + + async def feed(self, chunk: bytes) -> bytes: + complete_frames, self._pending = split_complete_sse_frames(self._pending + chunk) + if not complete_frames: + return b"" + if self._replacement is None: + self._replacement = await self._mint(complete_frames) + return self._rewrite(complete_frames) + + def flush(self) -> bytes: + tail: Final = self._pending + self._pending = b"" + return self._rewrite(tail) + + async def _mint(self, frames: bytes) -> tuple[bytes, bytes] | None: + first: Final = _first_streamed_response(frames) + if first is None: + return None + raw_id, snapshot = first + managed_id: Final = await _mint_or_reuse_object( + raw_id, + self._provider, + "response", + snapshot, + self._user_api_key_dict, + self._prisma_client, + self._is_create_route, + ) + return raw_id.encode(), managed_id.encode() + + def _rewrite(self, frames: bytes) -> bytes: + if self._replacement is None: + return frames + raw_id, managed_id = self._replacement + return frames.replace(raw_id, managed_id) + + +async def rewrite_streamed_response_ids( + stream: AsyncGenerator[bytes, None], + provider: str, + method: str, + route: str, + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient, +) -> AsyncGenerator[bytes, None]: + """ + Record ownership of the response object streamed back by a Responses API + passthrough and swap its managed id into every SSE frame, so a streamed + response is owned and resolved exactly like a non-streamed one. + + Streams for any other ``(provider, method, route)`` are relayed untouched. + """ + from litellm.proxy.auth.auth_utils import normalize_request_route + + canonical: Final = normalize_request_route(_canonical_path(route)) + field_specs: Final = BUILTIN_OUTPUT_ID_FIELD_MAP.get((provider, method, canonical), ()) + if _STREAMED_RESPONSE_ID_SPEC not in field_specs: + async for chunk in stream: + yield chunk + return + + rewriter: Final = _StreamedResponseIdRewriter( + provider=provider, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + is_create_route="{" not in canonical, + ) + async for chunk in stream: + rewritten_frames = await rewriter.feed(chunk) + if rewritten_frames: + yield rewritten_frames + tail: Final = rewriter.flush() + if tail: + yield tail + + # --------------------------------------------------------------------------- # List-route interception — serve listing entirely from DB # --------------------------------------------------------------------------- diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 60b85cb42d0..3d60f4f5f3a 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -1209,14 +1209,19 @@ async def pass_through_request( return StreamingResponse( wrap_passthrough_sse_bytes_with_keepalive_pings( - stream=PassThroughStreamingHandler.chunk_processor( - response=response, - request_body=_parsed_body, - litellm_logging_obj=logging_obj, - endpoint_type=endpoint_type, - start_time=start_time, - passthrough_success_handler_obj=pass_through_endpoint_logging, - url_route=str(url), + stream=_own_streamed_managed_ids( + stream=PassThroughStreamingHandler.chunk_processor( + response=response, + request_body=_parsed_body, + litellm_logging_obj=logging_obj, + endpoint_type=endpoint_type, + start_time=start_time, + passthrough_success_handler_obj=pass_through_endpoint_logging, + url_route=str(url), + ), + managed_id_provider=_managed_id_provider, + request=request, + user_api_key_dict=user_api_key_dict, ), ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, upstream_headers=response.headers, @@ -1285,14 +1290,19 @@ async def pass_through_request( return StreamingResponse( wrap_passthrough_sse_bytes_with_keepalive_pings( - stream=PassThroughStreamingHandler.chunk_processor( - response=response, - request_body=_parsed_body, - litellm_logging_obj=logging_obj, - endpoint_type=endpoint_type, - start_time=start_time, - passthrough_success_handler_obj=pass_through_endpoint_logging, - url_route=str(url), + stream=_own_streamed_managed_ids( + stream=PassThroughStreamingHandler.chunk_processor( + response=response, + request_body=_parsed_body, + litellm_logging_obj=logging_obj, + endpoint_type=endpoint_type, + start_time=start_time, + passthrough_success_handler_obj=pass_through_endpoint_logging, + url_route=str(url), + ), + managed_id_provider=_managed_id_provider, + request=request, + user_api_key_dict=user_api_key_dict, ), ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, upstream_headers=response.headers, @@ -2441,6 +2451,36 @@ def _is_streaming_response(response: httpx.Response) -> bool: return False +def _own_streamed_managed_ids( + stream: AsyncGenerator[bytes, None], + managed_id_provider: str | None, + request: Request, + user_api_key_dict: UserAPIKeyAuth, +) -> AsyncGenerator[bytes, None]: + from litellm.proxy.proxy_server import general_settings, prisma_client, proxy_logging_obj + + if ( + managed_id_provider is None + or not general_settings.get("passthrough_managed_object_ids", False) + or prisma_client is None + or proxy_logging_obj.get_proxy_hook("managed_files") is None + ): + return stream + from litellm.proxy.auth.auth_utils import get_request_route + from litellm.proxy.pass_through_endpoints.managed_id_rewriter import ( + rewrite_streamed_response_ids, + ) + + return rewrite_streamed_response_ids( + stream=stream, + provider=managed_id_provider, + method=request.method, + route=get_request_route(request), + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + + def _should_buffer_passthrough_response(response: httpx.Response) -> bool: """ Decide from the response headers whether the body must be read into memory. diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index b71622fc33d..eea1b19dea3 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -10,6 +10,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.proxy._types import PassThroughEndpointLoggingResultValues from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType from litellm.types.utils import StandardPassThroughResponseObject @@ -101,7 +102,7 @@ class PassThroughStreamingHandler: async for chunk in response.aiter_bytes(): raw_bytes.append(chunk) PassThroughStreamingHandler._stamp_first_chunk_if_needed(litellm_logging_obj) - complete_frames, pending = PassThroughStreamingHandler._split_complete_sse_frames( + complete_frames, pending = split_complete_sse_frames( pending + chunk ) # rebind-ok: SSE frame reassembly buffer across transport chunks if complete_frames: @@ -139,17 +140,6 @@ class PassThroughStreamingHandler: except Exception as e: verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e) - @staticmethod - def _split_complete_sse_frames(pending: bytes) -> tuple[bytes, bytes]: - lf_boundary_end: Final = pending.rfind(b"\n\n") + 2 - crlf_boundary_end: Final = pending.rfind(b"\r\n\r\n") + 4 - boundary_end: Final = max( - lf_boundary_end if lf_boundary_end >= 2 else 0, crlf_boundary_end if crlf_boundary_end >= 4 else 0 - ) - if boundary_end == 0: - return b"", pending - return pending[:boundary_end], pending[boundary_end:] - @staticmethod async def _route_streaming_logging_to_handler( litellm_logging_obj: LiteLLMLoggingObj, diff --git a/litellm/router.py b/litellm/router.py index de60d46e01d..1e9b23b2fa1 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7344,7 +7344,7 @@ class Router: ): raise error # then raise the error - if isinstance(error, openai.AuthenticationError): + if isinstance(error, (openai.AuthenticationError, openai.PermissionDeniedError)): """ - if other deployments available -> retry - else -> raise error diff --git a/tests/e2e/coverage_registry/README.md b/tests/e2e/coverage_registry/README.md index 5627c88dee4..da6aee84cc4 100644 --- a/tests/e2e/coverage_registry/README.md +++ b/tests/e2e/coverage_registry/README.md @@ -77,6 +77,24 @@ Strict mode exits non-zero on `@pytest.mark.covers(...)` ids that are not checke the registry. Add `--fail-on-collection-errors` when the job should also fail on pytest collection errors. +## Provider x feature matrix: customer-run Bedrock combinations + +The provider and feature combinations customers actually run get explicit cells, expanded +here as incidents surface new ones. The current Bedrock set, seeded from a customer's +production shape (regional `us.anthropic.*` inference-profile ids over both chat routes, +provider response headers for AWS-side correlation, and the Test Connection probe for a +responses-mode Bedrock Mantle deployment): + +| Cell | Feature | Covering test | +|------|---------|---------------| +| `llm.chat_completions.bedrock_converse.basic.nonstream.works` | regional `us.` id, Converse | `llm_translation/test_chat_completions_regression_e2e.py` | +| `llm.chat_completions.bedrock_converse.basic.stream.works` | regional `us.` id, Converse stream | `llm_translation/test_chat_completions_regression_e2e.py` | +| `llm.chat_completions.bedrock_invoke.basic.nonstream.works` | regional `us.` id, Invoke | `llm_translation/test_bedrock_provider_matrix_e2e.py` | +| `llm.chat_completions.bedrock_invoke.basic.stream.works` | regional `us.` id, Invoke stream | `llm_translation/test_bedrock_provider_matrix_e2e.py` | +| `llm.chat_completions.bedrock_converse.response_headers.nonstream.works` | `llm_provider-*` headers | `llm_translation/test_bedrock_provider_matrix_e2e.py` | +| `llm.chat_completions.bedrock_converse.response_headers.stream.works` | `llm_provider-*` headers, stream | `llm_translation/test_bedrock_provider_matrix_e2e.py` | +| `mgmt.model.test_connection.happy_path` | Test Connection, Bedrock Mantle | `management/test_model_test_connection_e2e.py` | + ## Status: this is a draft for review The cells were enumerated from the codebase and the tiers are a first proposal. Known diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index fc3803e5790..5662bdadb9c 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -29,6 +29,10 @@ - {id: llm.chat_completions.bedrock_converse.vision.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_converse, capability: vision, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Bedrock vision (Anthropic/Nova)"} - {id: llm.chat_completions.bedrock_converse.prompt_cache_5m.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: bedrock_converse, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Anthropic-on-Bedrock caching"} - {id: llm.chat_completions.bedrock_converse.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: bedrock_converse, capability: thinking, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Anthropic thinking on Bedrock"} +- {id: llm.chat_completions.bedrock_converse.response_headers.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_converse, capability: response_headers, streaming: nonstream, assertions: [works], source: "llms/bedrock/chat/converse_handler.py:248", rationale: "Bedrock request ids must surface as llm_provider-* response headers on /chat/completions so callers can correlate calls with AWS-side logs (#37003)", fail_before_fix: proven} +- {id: llm.chat_completions.bedrock_converse.response_headers.stream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_converse, capability: response_headers, streaming: stream, assertions: [works], source: "llms/bedrock/chat/converse_handler.py:154", rationale: "The llm_provider-* headers must also surface on streaming /chat/completions, where CustomStreamWrapper carries them instead of the nonstream setter"} +- {id: llm.chat_completions.bedrock_invoke.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_invoke, capability: basic, streaming: nonstream, assertions: [works], source: "proxy_server.py:8455", rationale: "Regional inference-profile ids (us.anthropic.*) over the invoke route, the deployment shape behind a customer timeout report on v1.90.0"} +- {id: llm.chat_completions.bedrock_invoke.basic.stream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_invoke, capability: basic, streaming: stream, assertions: [works], source: "proxy_server.py:8455", rationale: "Streaming with regional inference-profile ids over the invoke route"} - {id: llm.chat_completions.vertex.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "proxy_server.py:8455", rationale: "P0 route; Vertex AI"} - {id: llm.chat_completions.gemini.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: gemini, capability: basic, streaming: nonstream, assertions: [works], source: "test_chat_completions_regression_e2e.py", rationale: "Gemini OpenAI-compatible chat translation"} - {id: llm.chat_completions.gemini.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: chat_completions, route: gemini, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "test_chat_completions_regression_e2e.py", rationale: "Gemini chat cost lands in SpendLogs"} diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index d8788d7fcb0..1e6de0c3d6a 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -75,3 +75,4 @@ - {id: mgmt.workflow.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "workflow_management_endpoints.py", rationale: "Workflow tracking (smoke)"} - {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"} - {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"} +- {id: mgmt.model.test_connection.happy_path, module: mgmt, tier: P0, surface: api, assertions: [happy_path], source: "_health_endpoints.py:1785", rationale: "Test Connection for a responses-mode Bedrock Mantle deployment reaches the live provider and reports success; this exact shape 500ed on an acompletion partial before v1.91.0", fail_before_fix: proven} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index a5c723f8965..03d15f532b8 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -71,6 +71,7 @@ LlmCapability = Literal[ "pdf_input", "prompt_cache_1h", "prompt_cache_5m", + "response_headers", "service_tier", "structured_output", "thinking", diff --git a/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py b/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py new file mode 100644 index 00000000000..3c6aaa75ab3 --- /dev/null +++ b/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py @@ -0,0 +1,157 @@ +"""Live e2e for the Bedrock cells of the provider-feature matrix: provider +response headers on /chat/completions and regional inference-profile model ids +(us.anthropic.*) over the invoke route. + +Header forwarding is the #37003 contract: the proxy surfaces Bedrock's response +headers prefixed llm_provider- (llm_provider-x-amzn-requestid above all) so a +caller can hand AWS support the request id behind a completion. Regional +inference-profile ids are the deployment shape most Bedrock customers run; a +v1.90.0 regression timed them out, and the Converse route keeps them covered in +test_chat_completions_regression_e2e.py, so the invoke route carries its own +rows here. +""" + +from __future__ import annotations + +import pytest +from pydantic import BaseModel + +from e2e_config import unique_marker +from e2e_http import StreamingResponse, unwrap +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody +from passthrough_client import PassthroughClient + +pytestmark = pytest.mark.e2e + +CONVERSE_REGIONAL_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" +INVOKE_REGIONAL_BACKEND = "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0" +PROVIDER_HEADER_PREFIX = "llm_provider-" +BEDROCK_REQUEST_ID_HEADER = "llm_provider-x-amzn-requestid" + + +class _StreamDelta(BaseModel): + content: str | None = None + + +class _StreamChoice(BaseModel): + delta: _StreamDelta = _StreamDelta() + + +class _StreamChunk(BaseModel): + choices: list[_StreamChoice] = [] + + +def _streamed_text(events: list[str]) -> str: + chunks = [_StreamChunk.model_validate_json(event) for event in events] + return "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) + + +def _assert_streamed_completion(result: StreamingResponse) -> None: + assert result.ok and result.is_streaming, f"stream was not established: {result}" + assert result.stream_error is None, f"stream carried an error event: {result.stream_error}" + assert len(result.stream_events) > 1, f"stream did not deliver multiple data events: {result}" + assert _streamed_text(result.stream_events).strip(), ( + f"stream completed with no content deltas: {result.stream_events[:3]}" + ) + + +def _assert_request_id_header(result: StreamingResponse) -> None: + forwarded = [name for name in result.headers if name.startswith(PROVIDER_HEADER_PREFIX)] + assert result.headers.get(BEDROCK_REQUEST_ID_HEADER), ( + f"missing {BEDROCK_REQUEST_ID_HEADER}; forwarded provider headers: {forwarded}" + ) + + +def _assert_completion(response: ChatResponse) -> None: + assert response.choices, f"completion returned no choices: {response}" + message = response.choices[0].message + content = (message.content if message else None) or "" + assert content.strip(), f"completion carried no content: {response}" + + +def _register_bedrock_model( + client: PassthroughClient, resources: ResourceManager, prefix: str, backend: str +) -> str: + model = f"{prefix}-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody( + model=backend, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model + + +def _prompt() -> list[ChatMessage]: + return [ChatMessage(role="user", content="reply with one word")] + + +class TestBedrockResponseHeaders: + @pytest.mark.covers( + "llm.chat_completions.bedrock_converse.response_headers.nonstream.works", + exercised_on=[], + ) + def test_bedrock_request_id_header_surfaces( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = _register_bedrock_model(client, resources, "e2e-bedrock-headers", CONVERSE_REGIONAL_BACKEND) + key = resources.key() + + result = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(key), + json=ChatBody(model=model, messages=_prompt(), max_tokens=64), + ) + + assert result.ok, f"chat call failed: {result.status_code} {result.body[:300]}" + _assert_request_id_header(result) + + @pytest.mark.covers( + "llm.chat_completions.bedrock_converse.response_headers.stream.works", + exercised_on=[], + ) + def test_bedrock_request_id_header_surfaces_on_stream( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = _register_bedrock_model( + client, resources, "e2e-bedrock-headers-stream", CONVERSE_REGIONAL_BACKEND + ) + key = resources.key() + + result = client.proxy.chat_stream( + key, ChatBody(model=model, messages=_prompt(), stream=True, max_tokens=64) + ) + + _assert_streamed_completion(result) + _assert_request_id_header(result) + + +class TestBedrockInvokeRegionalModelIds: + @pytest.mark.covers("llm.chat_completions.bedrock_invoke.basic.nonstream.works", exercised_on=[]) + def test_invoke_regional_id_completes( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = _register_bedrock_model(client, resources, "e2e-bedrock-invoke", INVOKE_REGIONAL_BACKEND) + key = resources.key() + + response = unwrap(client.proxy.chat(key, ChatBody(model=model, messages=_prompt(), max_tokens=64))) + + _assert_completion(response) + + @pytest.mark.covers("llm.chat_completions.bedrock_invoke.basic.stream.works", exercised_on=[]) + def test_invoke_regional_id_streams( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = _register_bedrock_model(client, resources, "e2e-bedrock-invoke-stream", INVOKE_REGIONAL_BACKEND) + key = resources.key() + + result = client.proxy.chat_stream( + key, ChatBody(model=model, messages=_prompt(), stream=True, max_tokens=64) + ) + + _assert_streamed_completion(result) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index cdc31aeea79..b2bd41e19ba 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -14,6 +14,8 @@ from e2e_http import NoBody, ProbeResult, Result, StreamingResponse, Success, Un from models import ( ChatBody, ChatMessage, + ConnectionTestBody, + ConnectionTestResponse, CustomerDeleteBody, CustomerInfoParams, CustomerNewBody, @@ -118,6 +120,17 @@ class ManagementClient: ) ) + def connection_test(self, body: ConnectionTestBody) -> Result[ConnectionTestResponse]: + """POST /health/test_connection, the call behind the Admin UI's Test + Connection button, probing the live provider with the supplied params.""" + return self.proxy.transport.post( + "/health/test_connection", + headers=self.proxy.transport.master, + json=body, + response_type=ConnectionTestResponse, + timeout=120.0, + ) + def block_key(self, key: str) -> None: _ = unwrap( self.proxy.transport.post( diff --git a/tests/e2e/management/test_model_test_connection_e2e.py b/tests/e2e/management/test_model_test_connection_e2e.py new file mode 100644 index 00000000000..25b0b4f24e6 --- /dev/null +++ b/tests/e2e/management/test_model_test_connection_e2e.py @@ -0,0 +1,67 @@ +"""Live e2e for POST /health/test_connection, the API behind the Admin UI's +Test Connection button on the add-model form. + +The covered cell is a responses-mode Bedrock Mantle deployment: exactly this +shape 500ed on a functools.partial acompletion conflict before v1.91.0 while +every chat-mode probe stayed green, so the happy path asserts a real success +verdict from the live provider rather than just a 200 envelope. The region is a +literal because the endpoint rejects request-supplied os.environ/ references; +credentials fall through to the proxy's own environment (bearer token locally, +pod identity in CI). + +The endpoint caps every probe at HEALTH_CHECK_TIMEOUT_SECONDS and answers a +timed-out probe with HTTP 200 and an in-body "Timeout exceeded", which the +harness's status-code retry policy cannot see. A Mantle probe can hit that cap +transiently while the rest of the suite saturates the same AWS account, so only +that exact error is retried here; any other error verdict fails immediately. +""" + +from __future__ import annotations + +import time + +import pytest + +from e2e_http import unwrap +from management_client import ManagementClient +from models import ConnectionTestBody, ConnectionTestResponse, LiteLLMParamsBody + +pytestmark = pytest.mark.e2e + +MANTLE_RESPONSES_BACKEND = "bedrock_mantle/openai.gpt-5.6-luna" +MANTLE_REGION = "us-east-1" +PROBE_TIMEOUT_ERROR = "Timeout exceeded" +PROBE_ATTEMPTS = 3 +PROBE_RETRY_SLEEP_SECONDS = 30 + + +def _probe_mantle(client: ManagementClient) -> ConnectionTestResponse: + return unwrap( + client.connection_test( + ConnectionTestBody( + litellm_params=LiteLLMParamsBody( + model=MANTLE_RESPONSES_BACKEND, aws_region_name=MANTLE_REGION + ), + mode="responses", + ) + ) + ) + + +class TestModelTestConnection: + @pytest.mark.covers("mgmt.model.test_connection.happy_path") + def test_bedrock_mantle_responses_connection_succeeds(self, client: ManagementClient) -> None: + for attempt in range(1, PROBE_ATTEMPTS + 1): + response = _probe_mantle(client) + if response.status == "success": + return + error = response.result.error if response.result else None + assert error == PROBE_TIMEOUT_ERROR, f"test_connection reported an error: {error}" + if attempt < PROBE_ATTEMPTS: + print( + f"test_connection probe timed out; retry {attempt}/{PROBE_ATTEMPTS - 1}" + f" in {PROBE_RETRY_SLEEP_SECONDS}s", + flush=True, + ) + time.sleep(PROBE_RETRY_SLEEP_SECONDS) + pytest.fail(f"test_connection timed out on all {PROBE_ATTEMPTS} attempts") diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 77555a9fc95..95a02b58824 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -857,6 +857,26 @@ class ModelDeleteBody(BaseModel): id: str +class ConnectionTestBody(BaseModel): + """POST /health/test_connection body, the API behind the Admin UI's Test + Connection button: the deployment params as typed into the add-model form and + the health-check mode picking which endpoint the probe calls. The endpoint + rejects `os.environ/` references, so credentials are either literal values or + omitted to fall through to the proxy's own environment.""" + + litellm_params: LiteLLMParamsBody + mode: Literal["chat", "completion", "embedding", "responses"] + + +class ConnectionTestResult(BaseModel): + error: str | None = None + + +class ConnectionTestResponse(BaseModel): + status: Literal["success", "error"] + result: ConnectionTestResult | None = None + + class CredentialCreateBody(BaseModel): credential_name: str credential_values: dict[str, str] diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 6f513ce1bd4..c875bf5b535 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -716,6 +716,136 @@ def test_generic_cost_per_token_tier_without_an_output_rate_bills_the_model_rate litellm.model_cost.pop(model, None) +def test_generic_cost_per_token_tier_without_cache_rates_bills_cache_at_the_tier_input_rate(): + model = "litellm-test-tiered-no-cache-rates" + custom_llm_provider = "openrouter" + litellm.register_model( + { + model: { + "litellm_provider": custom_llm_provider, + "mode": "chat", + "cache_read_input_token_cost": 9e-09, + "cache_creation_input_token_cost": 9e-06, + "tiered_pricing": [ + { + "range": [0, 32000], + "input_cost_per_token": 4.6e-07, + "output_cost_per_token": 2.3e-06, + }, + { + "range": [32000, 128000], + "input_cost_per_token": 7e-07, + "output_cost_per_token": 3.5e-06, + }, + ], + } + } + ) + + try: + uncached = Usage(prompt_tokens=40000, completion_tokens=100, total_tokens=40100) + cached = Usage( + prompt_tokens=40000, + completion_tokens=100, + total_tokens=40100, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=5000, cache_creation_tokens=15000 + ), + ) + uncached_prompt_cost, _ = generic_cost_per_token( + model=model, + usage=uncached, + custom_llm_provider=custom_llm_provider, + ) + cached_prompt_cost, cached_completion_cost = generic_cost_per_token( + model=model, + usage=cached, + custom_llm_provider=custom_llm_provider, + ) + + tier_input_rate = 7e-07 + assert round(cached_prompt_cost, 12) == round(40000 * tier_input_rate, 12) + assert round(cached_prompt_cost, 12) == round(uncached_prompt_cost, 12) + assert round(cached_completion_cost, 12) == round(100 * 3.5e-06, 12) + finally: + litellm.model_cost.pop(model, None) + + +def test_generic_cost_per_token_tier_without_a_1hr_cache_rate_bills_the_tier_cache_creation_rate(): + model = "litellm-test-tiered-no-1hr-cache-rate" + custom_llm_provider = "openrouter" + litellm.register_model( + { + model: { + "litellm_provider": custom_llm_provider, + "mode": "chat", + "cache_creation_input_token_cost_above_1hr": 9e-05, + "tiered_pricing": [ + { + "range": [0, 128000], + "input_cost_per_token": 7e-07, + "output_cost_per_token": 3.5e-06, + "cache_creation_input_token_cost": 8.75e-07, + } + ], + } + } + ) + + try: + usage = Usage( + prompt_tokens=1000, + completion_tokens=10, + total_tokens=1010, + prompt_tokens_details=PromptTokensDetailsWrapper( + cache_creation_tokens=800, + cache_creation_token_details=CacheCreationTokenDetails( + ephemeral_5m_input_tokens=300, ephemeral_1h_input_tokens=500 + ), + ), + ) + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + ) + + tier_cache_creation_rate = 8.75e-07 + expected_prompt = (200 * 7e-07) + (800 * tier_cache_creation_rate) + assert round(prompt_cost, 12) == round(expected_prompt, 12) + assert round(completion_cost, 12) == round(10 * 3.5e-06, 12) + finally: + litellm.model_cost.pop(model, None) + + +def test_generic_cost_per_token_tier_without_an_input_rate_is_not_a_priced_tier(): + model = "litellm-test-tiered-no-input-rate" + custom_llm_provider = "openrouter" + litellm.register_model( + { + model: { + "litellm_provider": custom_llm_provider, + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + "tiered_pricing": [{"range": [0, 128000], "output_cost_per_token": 3.5e-06}], + } + } + ) + + try: + usage = Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100) + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + ) + assert round(prompt_cost, 12) == round(1000 * 1e-06, 12) + assert round(completion_cost, 12) == round(100 * 2e-06, 12) + finally: + litellm.model_cost.pop(model, None) + + def test_router_deployment_with_input_only_tiers_bills_completions_at_the_backend_rate(): """Regression: the router registers a deployment's custom pricing as a standalone model_cost entry holding only the supplied fields, so an input-only tier table left diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index d5d5004dbe8..895044c8ad5 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -13,6 +13,7 @@ from litellm.litellm_core_utils.exception_mapping_utils import ( extract_and_raise_litellm_exception, ) from litellm.llms.openai.common_utils import OpenAIError +from litellm.types.utils import LlmProviders # Test cases for is_error_str_context_window_exceeded # Tuple format: (error_message, expected_result) @@ -785,36 +786,24 @@ OPENAI_SHAPED = { 503: (litellm.ServiceUnavailableError, 503), } -UPSTREAM_STATUS_DISCARDED = (litellm.APIConnectionError, 500) +PERMISSION_DENIED = (litellm.PermissionDeniedError, 403) -PROVIDERS_THAT_DISCARD_THE_UPSTREAM_STATUS = ("cloudflare", "ollama", "vllm") +STATUS_KEYED = {**OPENAI_SHAPED, 403: PERMISSION_DENIED} DEVIATIONS_FROM_THE_OPENAI_SHAPE = { - "anthropic": { - 403: (litellm.PermissionDeniedError, 403), - 422: UPSTREAM_STATUS_DISCARDED, - }, + "anthropic": {403: PERMISSION_DENIED}, "azure": {500: (litellm.APIError, 500)}, "bedrock": { - 403: UPSTREAM_STATUS_DISCARDED, + 403: PERMISSION_DENIED, 500: (litellm.ServiceUnavailableError, 503), }, - "cohere": { - 401: UPSTREAM_STATUS_DISCARDED, - 403: UPSTREAM_STATUS_DISCARDED, - 404: UPSTREAM_STATUS_DISCARDED, - 422: UPSTREAM_STATUS_DISCARDED, - 429: UPSTREAM_STATUS_DISCARDED, - 503: UPSTREAM_STATUS_DISCARDED, - }, + "cloudflare": {403: PERMISSION_DENIED}, + "cohere": {403: PERMISSION_DENIED}, "databricks": { - 403: (litellm.AuthenticationError, 401), + 403: PERMISSION_DENIED, 422: (litellm.BadRequestError, 400), }, - "gemini": { - 403: (litellm.PermissionDeniedError, 403), - 422: UPSTREAM_STATUS_DISCARDED, - }, + "gemini": {403: PERMISSION_DENIED}, "huggingface": { 404: (litellm.APIError, 404), 422: (litellm.APIError, 422), @@ -827,6 +816,7 @@ DEVIATIONS_FROM_THE_OPENAI_SHAPE = { 500: (litellm.APIError, 500), 503: (litellm.APIError, 503), }, + "ollama": {403: PERMISSION_DENIED}, "openrouter": {500: (litellm.APIError, 500)}, "replicate": { 403: (litellm.APIError, 500), @@ -836,17 +826,11 @@ DEVIATIONS_FROM_THE_OPENAI_SHAPE = { 503: (litellm.APIError, 500), }, "sagemaker": { - 403: UPSTREAM_STATUS_DISCARDED, + 403: PERMISSION_DENIED, 500: (litellm.ServiceUnavailableError, 503), }, - "vertex_ai": { - 403: (litellm.PermissionDeniedError, 403), - 422: UPSTREAM_STATUS_DISCARDED, - }, - **{ - provider: dict.fromkeys(UPSTREAM_STATUS_CODES, UPSTREAM_STATUS_DISCARDED) - for provider in PROVIDERS_THAT_DISCARD_THE_UPSTREAM_STATUS - }, + "vertex_ai": {403: PERMISSION_DENIED}, + "vllm": {403: PERMISSION_DENIED}, } PROVIDERS_WITH_A_HANDLER = ( @@ -878,6 +862,38 @@ PROVIDERS_WITH_A_HANDLER = ( "xai", ) +PROVIDER_ALIASES_WITH_A_HANDLER = ( + "aleph_alpha", + "anthropic_text", + "azure_text", + "bedrock_mantle", + "cohere_chat", + "custom_openai", + "lemonade", + "litellm_proxy", + "ollama_chat", + "predibase", + "sagemaker_chat", + "text-completion-openai", + "vertex_ai_beta", + "watsonx", +) + +PROVIDERS_WITHOUT_A_HANDLER = tuple( + sorted( + frozenset(provider.value for provider in LlmProviders) + - frozenset(PROVIDERS_WITH_A_HANDLER) + - frozenset(PROVIDER_ALIASES_WITH_A_HANDLER) + - frozenset(litellm.openai_compatible_providers) + ) +) + +MINIMAX_401_BODY = ( + '{"type":"error","error":{"type":"authorized_error","message":"login fail: Please carry the API secret key ' + "in the 'Authorization' field of the request header (1004)\",\"http_code\":\"401\"}," + '"request_id":"06ddc9ba97ee6340e38f10e09787f547"}' +) + def _expected_for(provider: str, status_code: int) -> tuple[type[Exception], int]: return DEVIATIONS_FROM_THE_OPENAI_SHAPE.get(provider, {}).get( @@ -941,6 +957,51 @@ def test_an_already_mapped_litellm_exception_passes_through_untouched( assert returned is already_mapped +@pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES) +@pytest.mark.parametrize("provider", PROVIDERS_WITHOUT_A_HANDLER) +def test_a_provider_without_a_handler_maps_by_the_upstream_status( + provider, status_code, quiet_exception_mapping +): + expected_class, expected_status = STATUS_KEYED[status_code] + + with pytest.raises(openai.APIError) as raised: + exception_type( + model="test-model", + original_exception=_UpstreamHTTPError(status_code=status_code), + custom_llm_provider=provider, + ) + + assert type(raised.value) is expected_class + assert raised.value.status_code == expected_status + assert raised.value.llm_provider == provider + assert raised.value.model == "test-model" + + +def test_a_minimax_bad_key_is_an_authentication_error(quiet_exception_mapping): + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + with pytest.raises(litellm.AuthenticationError) as raised: + exception_type( + model="MiniMax-M2.5", + original_exception=BaseLLMException(status_code=401, message=MINIMAX_401_BODY), + custom_llm_provider="minimax", + ) + + assert raised.value.status_code == 401 + assert raised.value.llm_provider == "minimax" + assert raised.value.message.startswith("litellm.AuthenticationError: MinimaxException - ") + assert "login fail" in raised.value.message + + +def test_an_exception_without_a_status_is_still_a_connection_error(quiet_exception_mapping): + with pytest.raises(litellm.APIConnectionError): + exception_type( + model="MiniMax-M2.5", + original_exception=RuntimeError("socket hung up"), + custom_llm_provider="minimax", + ) + + CONTEXT_WINDOW_MESSAGE = "This model's maximum context length is 4096 tokens." CONTENT_POLICY_MESSAGE = ( '{"error": {"type": "invalid_request_error", "code": "content_policy_violation"}}' @@ -996,9 +1057,7 @@ class _UpstreamErrorWithMessage(_UpstreamHTTPError): def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it( provider, quiet_exception_mapping ): - if provider in PROVIDERS_THAT_DISCARD_THE_UPSTREAM_STATUS: - expected_class, expected_status = UPSTREAM_STATUS_DISCARDED - elif provider in PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW: + if provider in PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW: expected_class, expected_status = litellm.ContextWindowExceededError, 400 else: expected_class, expected_status = litellm.BadRequestError, 400 @@ -1018,9 +1077,7 @@ def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it( def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it( provider, quiet_exception_mapping ): - if provider in PROVIDERS_THAT_DISCARD_THE_UPSTREAM_STATUS: - expected_class, expected_status = UPSTREAM_STATUS_DISCARDED - elif provider in PROVIDERS_THAT_RECOGNISE_A_CONTENT_POLICY_BLOCK: + if provider in PROVIDERS_THAT_RECOGNISE_A_CONTENT_POLICY_BLOCK: expected_class, expected_status = litellm.ContentPolicyViolationError, 400 else: expected_class, expected_status = litellm.BadRequestError, 400 diff --git a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py index e6aad7688d1..35a54332f66 100644 --- a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py +++ b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py @@ -322,7 +322,7 @@ async def test_query_param_key_not_leaked_with_dummy_caller_key( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", fake_get, ): - with pytest.raises(litellm.APIConnectionError): + with pytest.raises(litellm.InternalServerError): await litellm.asearch( query="secrets", search_provider=provider, diff --git a/tests/test_litellm/llms/compactifai/test_compactifai.py b/tests/test_litellm/llms/compactifai/test_compactifai.py index fef0baf2884..fd31049731a 100644 --- a/tests/test_litellm/llms/compactifai/test_compactifai.py +++ b/tests/test_litellm/llms/compactifai/test_compactifai.py @@ -172,7 +172,7 @@ def test_compactifai_authentication_error(respx_mock): json=mock_error, status_code=401 ) - with pytest.raises(litellm.APIConnectionError) as exc_info: + with pytest.raises(litellm.AuthenticationError) as exc_info: litellm.completion( model="compactifai/cai-llama-3-1-8b-slim", messages=[{"role": "user", "content": "test"}], diff --git a/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py b/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py index 0c241add77b..383a7afbe93 100644 --- a/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py +++ b/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py @@ -233,7 +233,7 @@ def test_langflow_extra_body_cannot_inject_tweaks_into_run_payload(): return resp with patch.object(HTTPHandler, "post", side_effect=fake_post): - with pytest.raises(litellm.APIConnectionError): + with pytest.raises(litellm.BadRequestError): litellm.completion( model="langflow/my-flow", messages=[{"role": "user", "content": "hello"}], diff --git a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py index e3a958ff167..bee8ab3a9ac 100644 --- a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py +++ b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py @@ -20,11 +20,24 @@ TOOL_CALLING_MODEL = "openai/gpt-oss-20b" REASONING_MODEL = "deepseek-ai/DeepSeek-V3.1" UNMAPPED_MODEL = "example-org/brand-new-model" NO_TOOLS_MODEL = "example-org/no-tools-model" +NO_SCHEMA_MODEL = "example-org/no-schema-model" TOOL_PARAMS = ("tools", "tool_choice", "function_call") WEATHER_TOOLS = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}] +VOICE_NOTE_SCHEMA = { + "type": "object", + "properties": {"title": {"type": "string"}, "summary": {"type": "string"}}, + "required": ["title", "summary"], + "additionalProperties": False, +} +JSON_SCHEMA_RESPONSE_FORMAT = { + "type": "json_schema", + "json_schema": {"name": "voice_note", "schema": VOICE_NOTE_SCHEMA, "strict": True}, +} +REGEX_RESPONSE_FORMAT = {"type": "regex", "pattern": "(positive|neutral|negative)"} + @pytest.fixture(autouse=True) def force_local_model_cost(monkeypatch): @@ -48,6 +61,15 @@ def registry_disables_function_calling(monkeypatch): ) +@pytest.fixture +def registry_disables_response_schema(monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + f"together_ai/{NO_SCHEMA_MODEL}", + {"litellm_provider": "together_ai", "mode": "chat", "supports_response_schema": False}, + ) + + @pytest.fixture def together_warning_log(caplog): from litellm._logging import verbose_logger @@ -70,7 +92,7 @@ def test_supported_params_unmapped_model_keeps_tool_params(): for param in TOOL_PARAMS: assert param in supported - assert "response_format" not in supported + assert "response_format" in supported assert "stream" in supported assert "temperature" in supported @@ -80,7 +102,7 @@ def test_supported_params_no_tools_model_keeps_tool_params(registry_disables_fun for param in TOOL_PARAMS: assert param in supported - assert "response_format" not in supported + assert "response_format" in supported def test_map_openai_params_tool_calling_model_passes_tools(): @@ -148,21 +170,17 @@ def test_map_openai_params_reasoning_model_passes_sampling_params(): assert mapped["max_tokens"] == 512 -def test_map_openai_params_drops_text_response_format(): - mapped = TogetherAIChatConfig().map_openai_params( - non_default_params={"response_format": {"type": "text"}, "temperature": 0.5}, - optional_params={}, - model=REASONING_MODEL, - drop_params=False, - ) - - assert "response_format" not in mapped - assert mapped["temperature"] == 0.5 - - -def test_map_openai_params_keeps_json_response_format(): - response_format = {"type": "json_object"} - +@pytest.mark.parametrize( + "response_format", + [ + {"type": "text"}, + {"type": "json_object"}, + {"type": "json_object", "schema": VOICE_NOTE_SCHEMA}, + JSON_SCHEMA_RESPONSE_FORMAT, + REGEX_RESPONSE_FORMAT, + ], +) +def test_map_openai_params_schema_model_passes_response_format_through(response_format): mapped = TogetherAIChatConfig().map_openai_params( non_default_params={"response_format": response_format}, optional_params={}, @@ -173,6 +191,46 @@ def test_map_openai_params_keeps_json_response_format(): assert mapped["response_format"] == response_format +@pytest.mark.parametrize("drop_params", [False, True]) +def test_map_openai_params_unmapped_model_passes_response_format_through(drop_params, together_warning_log): + mapped = TogetherAIChatConfig().map_openai_params( + non_default_params={"response_format": JSON_SCHEMA_RESPONSE_FORMAT}, + optional_params={}, + model=UNMAPPED_MODEL, + drop_params=drop_params, + ) + + assert mapped["response_format"] == JSON_SCHEMA_RESPONSE_FORMAT + assert UNMAPPED_MODEL in together_warning_log.text + assert "passing response_format through" in together_warning_log.text + + +def test_map_openai_params_no_schema_model_drops_response_format_with_warning( + registry_disables_response_schema, together_warning_log +): + mapped = TogetherAIChatConfig().map_openai_params( + non_default_params={"response_format": JSON_SCHEMA_RESPONSE_FORMAT, "temperature": 0.5}, + optional_params={}, + model=NO_SCHEMA_MODEL, + drop_params=True, + ) + + assert "response_format" not in mapped + assert mapped["temperature"] == 0.5 + assert NO_SCHEMA_MODEL in together_warning_log.text + assert "dropping response_format" in together_warning_log.text + + +def test_map_openai_params_no_schema_model_raises_without_drop_params(registry_disables_response_schema): + with pytest.raises(UnsupportedParamsError, match="response_format"): + TogetherAIChatConfig().map_openai_params( + non_default_params={"response_format": JSON_SCHEMA_RESPONSE_FORMAT}, + optional_params={}, + model=NO_SCHEMA_MODEL, + drop_params=False, + ) + + def _transform_response(message: dict) -> ModelResponse: raw_response_json = { "id": "chatcmpl-test", @@ -206,26 +264,20 @@ def _transform_response(message: dict) -> ModelResponse: def test_transform_response_maps_reasoning_to_reasoning_content(): - result = _transform_response( - {"role": "assistant", "content": "4", "reasoning": "2+2 equals 4"} - ) + result = _transform_response({"role": "assistant", "content": "4", "reasoning": "2+2 equals 4"}) assert result.choices[0].message.content == "4" assert result.choices[0].message.reasoning_content == "2+2 equals 4" def test_transform_response_preserves_reasoning_content_field(): - result = _transform_response( - {"role": "assistant", "content": "4", "reasoning_content": "adding 2 and 2"} - ) + result = _transform_response({"role": "assistant", "content": "4", "reasoning_content": "adding 2 and 2"}) assert result.choices[0].message.reasoning_content == "adding 2 and 2" def test_streaming_chunk_maps_delta_reasoning_to_reasoning_content(): - iterator = TogetherAIChatConfig().get_model_response_iterator( - streaming_response=iter(()), sync_stream=True - ) + iterator = TogetherAIChatConfig().get_model_response_iterator(streaming_response=iter(()), sync_stream=True) assert isinstance(iterator, OpenAIChatCompletionStreamingHandler) parsed = iterator.chunk_parser( @@ -241,9 +293,7 @@ def test_streaming_chunk_maps_delta_reasoning_to_reasoning_content(): def test_streaming_chunk_preserves_tool_call_index_and_id(): - iterator = TogetherAIChatConfig().get_model_response_iterator( - streaming_response=iter(()), sync_stream=True - ) + iterator = TogetherAIChatConfig().get_model_response_iterator(streaming_response=iter(()), sync_stream=True) def parse_tool_call_chunk(tool_call: dict): parsed = iterator.chunk_parser( @@ -374,9 +424,7 @@ def test_together_ai_config_alias_points_at_chat_config(): def test_provider_config_manager_returns_together_chat_config(): from litellm.utils import ProviderConfigManager - config = ProviderConfigManager.get_provider_chat_config( - model=REASONING_MODEL, provider=LlmProviders.TOGETHER_AI - ) + config = ProviderConfigManager.get_provider_chat_config(model=REASONING_MODEL, provider=LlmProviders.TOGETHER_AI) assert isinstance(config, TogetherAIChatConfig) @@ -484,6 +532,72 @@ def test_completion_unmapped_model_sends_tools_to_together(): assert json.loads(tool_call.function.arguments) == {"city": "San Francisco"} +def _capture_completion_request(model: str, **completion_kwargs) -> dict: + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + captured_requests = [] + + def respond(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-together-structured", + "object": "chat.completion", + "created": 1234567890, + "model": model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": '{"title": "t", "summary": "s"}'}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + ) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + litellm.completion( + model=f"together_ai/{model}", + messages=[{"role": "user", "content": "Summarize with a title and summary."}], + api_key="fake-key", + client=client, + **completion_kwargs, + ) + return json.loads(captured_requests[0].content) + + +def test_completion_unmapped_model_sends_json_schema_to_together(): + request_body = _capture_completion_request( + UNMAPPED_MODEL, response_format=JSON_SCHEMA_RESPONSE_FORMAT, drop_params=True + ) + + assert request_body["response_format"] == JSON_SCHEMA_RESPONSE_FORMAT + + +def test_completion_pydantic_response_format_sends_json_schema_to_together(): + from pydantic import BaseModel + + class VoiceNote(BaseModel): + title: str + summary: str + + request_body = _capture_completion_request(TOOL_CALLING_MODEL, response_format=VoiceNote) + + sent = request_body["response_format"] + assert sent["type"] == "json_schema" + assert sent["json_schema"]["name"] == "VoiceNote" + assert sent["json_schema"]["strict"] is True + assert sent["json_schema"]["schema"]["required"] == ["title", "summary"] + + +def test_completion_regex_response_format_sends_pattern_to_together(): + request_body = _capture_completion_request(TOOL_CALLING_MODEL, response_format=REGEX_RESPONSE_FORMAT) + + assert request_body["response_format"] == REGEX_RESPONSE_FORMAT + + TOGETHER_CHAT_URL = "https://api.together.ai/v1/chat/completions" WEATHER_AND_TIME_TOOLS = [ @@ -552,14 +666,23 @@ PARALLEL_TOOL_CALL_STREAM = ( _chunk( { "tool_calls": [ - {"index": 0, "id": "call_weather", "type": "function", "function": {"name": "get_weather", "arguments": ""}} + { + "index": 0, + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": ""}, + } ] } ), _chunk({"tool_calls": [{"index": 0, "function": {"arguments": '{"city": "San'}}]}), _chunk({"tool_calls": [{"index": 0, "function": {"arguments": ' Francisco"}'}}]}), _chunk( - {"tool_calls": [{"index": 1, "id": "call_time", "type": "function", "function": {"name": "get_time", "arguments": ""}}]} + { + "tool_calls": [ + {"index": 1, "id": "call_time", "type": "function", "function": {"name": "get_time", "arguments": ""}} + ] + } ), _chunk({"tool_calls": [{"index": 1, "function": {"arguments": '{"tz": "PST"}'}}]}, finish_reason="tool_calls"), ) @@ -802,10 +925,14 @@ def test_anthropic_messages_streams_together_tool_call_as_input_json_delta(): if event["type"] == "content_block_start" and event["content_block"]["type"] == "tool_use" } input_json_deltas = [ - event for event in events if event["type"] == "content_block_delta" and event["delta"]["type"] == "input_json_delta" + event + for event in events + if event["type"] == "content_block_delta" and event["delta"]["type"] == "input_json_delta" ] tool_inputs = { - block["name"]: json.loads("".join(delta["delta"]["partial_json"] for delta in input_json_deltas if delta["index"] == index)) + block["name"]: json.loads( + "".join(delta["delta"]["partial_json"] for delta in input_json_deltas if delta["index"] == index) + ) for index, block in tool_starts.items() } assert {block["id"] for block in tool_starts.values()} == {"call_weather", "call_time"} diff --git a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py index 39af9f08540..5f74f0f602f 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py @@ -238,7 +238,7 @@ class TestVertexGemmaCompletion: Expected: Proper error handling when 'predictions' field is missing """ - from litellm.exceptions import APIConnectionError + from litellm.exceptions import BadRequestError # Invalid response without predictions field invalid_response = { @@ -260,8 +260,8 @@ class TestVertexGemmaCompletion: mock_client.post = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client - # Should raise exception (wrapped as APIConnectionError by LiteLLM) - with pytest.raises(APIConnectionError) as exc_info: + # Should raise exception (wrapped as BadRequestError by LiteLLM) + with pytest.raises(BadRequestError) as exc_info: await litellm.acompletion( model="vertex_ai/gemma/gemma-3-12b-it", messages=[{"role": "user", "content": "Test"}], diff --git a/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py b/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py index 89ae74920fe..69b92f5e4d7 100644 --- a/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py +++ b/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py @@ -10,6 +10,7 @@ from litellm.proxy.common_utils.sse_keepalive import ( ANTHROPIC_PING_SSE_CHUNK, SSE_COMMENT_PING_BYTES, resolve_ttft_keepalive_interval, + split_complete_sse_frames, wrap_passthrough_sse_bytes_with_keepalive_pings, wrap_sse_stream_with_keepalive_pings, ) @@ -18,6 +19,19 @@ MESSAGE_START_CHUNK: Final = 'data: {"type": "message_start"}\n\n' TEXT_DELTA_CHUNK: Final = 'data: {"type": "content_block_delta"}\n\n' +@pytest.mark.parametrize("delimiter", [b"\n\n", b"\r\n\r\n", b"\r\r"]) +def test_split_complete_sse_frames_recognizes_every_sse_frame_delimiter(delimiter: bytes): + newline: Final = delimiter[: len(delimiter) // 2] + frame: Final = b"event: response.created" + newline + b"data: {}" + delimiter + tail: Final = b"data: partial" + + assert split_complete_sse_frames(frame + tail) == (frame, tail) + + +def test_split_complete_sse_frames_holds_bytes_with_no_complete_frame(): + assert split_complete_sse_frames(b"data: unterminated") == (b"", b"data: unterminated") + + @pytest.mark.asyncio async def test_pings_fill_mid_stream_silence_and_preserve_chunk_order(): async def gappy_stream() -> AsyncGenerator[str, None]: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_managed_id_rewriter.py b/tests/test_litellm/proxy/pass_through_endpoints/test_managed_id_rewriter.py index f5bec4a2585..dc8c49b93d8 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_managed_id_rewriter.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_managed_id_rewriter.py @@ -1,12 +1,15 @@ import datetime +import json +from collections.abc import AsyncIterator, Iterable from unittest.mock import AsyncMock, MagicMock import pytest from litellm.proxy._types import ProxyException, UserAPIKeyAuth -from litellm.proxy.pass_through_endpoints.managed_id_codec import new_managed_id +from litellm.proxy.pass_through_endpoints.managed_id_codec import decode, new_managed_id from litellm.proxy.pass_through_endpoints.managed_id_rewriter import ( list_passthrough_ids_from_db, + rewrite_streamed_response_ids, ) @@ -27,9 +30,39 @@ def _prisma_client(file_rows=None, batch_rows=None) -> MagicMock: pc.db.litellm_managedobjecttable.find_many = AsyncMock( side_effect=lambda *args, take=None, **kwargs: list(batch_rows or [])[:take] ) + pc.db.litellm_managedobjecttable.upsert = AsyncMock(return_value=None) return pc +RAW_RESPONSE_ID = "resp_0123456789abcdef" + + +def _response_stream_bytes(raw_id: str = RAW_RESPONSE_ID) -> bytes: + events = ( + ("response.created", {"type": "response.created", "response": {"id": raw_id, "status": "in_progress"}}), + ("response.output_text.delta", {"type": "response.output_text.delta", "delta": "mango"}), + ("response.completed", {"type": "response.completed", "response": {"id": raw_id, "status": "completed"}}), + ) + return b"".join(f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events) + + +async def _chunks(payload: bytes, size: int) -> AsyncIterator[bytes]: + for start in range(0, len(payload), size): + yield payload[start : start + size] + + +async def _collect(stream: AsyncIterator[bytes]) -> bytes: + return b"".join([chunk async for chunk in stream]) + + +def _response_ids(sse: bytes) -> Iterable[str]: + for line in sse.decode().splitlines(): + if line.startswith("data:"): + event = json.loads(line[len("data:") :]) + if "response" in event: + yield event["response"]["id"] + + def _file_row(unified_id: str) -> MagicMock: row = MagicMock() row.unified_file_id = unified_id @@ -67,9 +100,7 @@ def _batch_row(unified_id: str) -> MagicMock: ), ], ) -async def test_list_batches_out_of_range_limit_raises_400( - limit, expected_message, expected_openai_code -): +async def test_list_batches_out_of_range_limit_raises_400(limit, expected_message, expected_openai_code): pc = _prisma_client(batch_rows=[_batch_row(new_managed_id("openai", "batch_abc"))]) with pytest.raises(ProxyException) as exc: @@ -147,3 +178,98 @@ async def test_list_files_drops_batch_guardrail_key_persisted_by_an_older_proxy( assert result is not None assert "litellm_batch_guardrail" not in result["data"][0] assert result["data"][0]["filename"] == "test.jsonl" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chunk_size", [1, 7, 4096]) +async def test_streamed_response_is_owned_and_rewritten_across_chunk_boundaries(chunk_size: int): + """A streamed POST /v1/responses records the caller as owner once and returns + the minted id in every event, no matter how the transport splits the SSE bytes.""" + pc = _prisma_client() + + output = await _collect( + rewrite_streamed_response_ids( + stream=_chunks(_response_stream_bytes(), chunk_size), + provider="openai", + method="POST", + route="/openai_passthrough/v1/responses", + user_api_key_dict=_user(), + prisma_client=pc, + ) + ) + + pc.db.litellm_managedobjecttable.upsert.assert_awaited_once() + created = pc.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]["create"] + assert created["created_by"] == "user-1" + assert created["team_id"] == "team-1" + assert created["file_purpose"] == "response" + assert created["model_object_id"] == f"passthrough:openai:{RAW_RESPONSE_ID}" + managed_id = created["unified_object_id"] + assert decode(managed_id).raw_provider_id == RAW_RESPONSE_ID + assert list(_response_ids(output)) == [managed_id, managed_id] + assert RAW_RESPONSE_ID.encode() not in output + assert output == _response_stream_bytes(managed_id) + + +@pytest.mark.asyncio +async def test_streamed_response_with_cr_only_frame_delimiters_is_still_owned_and_rewritten(): + """SSE also terminates lines with a lone CR; those frames must mint and rewrite too.""" + pc = _prisma_client() + payload = _response_stream_bytes().replace(b"\n", b"\r") + + output = await _collect( + rewrite_streamed_response_ids( + stream=_chunks(payload, 7), + provider="openai", + method="POST", + route="/openai_passthrough/v1/responses", + user_api_key_dict=_user(), + prisma_client=pc, + ) + ) + + pc.db.litellm_managedobjecttable.upsert.assert_awaited_once() + managed_id = pc.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]["create"]["unified_object_id"] + assert RAW_RESPONSE_ID.encode() not in output + assert output == _response_stream_bytes(managed_id).replace(b"\n", b"\r") + + +@pytest.mark.asyncio +async def test_streamed_bytes_untouched_on_routes_without_a_response_id(): + pc = _prisma_client() + payload = _response_stream_bytes() + + output = await _collect( + rewrite_streamed_response_ids( + stream=_chunks(payload, 5), + provider="openai", + method="POST", + route="/openai_passthrough/v1/chat/completions", + user_api_key_dict=_user(), + prisma_client=pc, + ) + ) + + assert output == payload + pc.db.litellm_managedobjecttable.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_streamed_response_stays_raw_and_intact_when_the_row_cannot_be_persisted(): + pc = _prisma_client() + pc.db.litellm_managedobjecttable.upsert = AsyncMock(side_effect=RuntimeError("db down")) + payload = _response_stream_bytes() + + output = await _collect( + rewrite_streamed_response_ids( + stream=_chunks(payload, 3), + provider="openai", + method="POST", + route="/openai_passthrough/v1/responses", + user_api_key_dict=_user(), + prisma_client=pc, + ) + ) + + assert output == payload + pc.db.litellm_managedobjecttable.upsert.assert_awaited_once() diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 99a84d43c9b..a3f56adb86f 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -1493,6 +1493,86 @@ async def test_pass_through_request_sse_response_marks_logging_obj_as_stream(): assert logging_obj.model_call_details["stream"] is True +@pytest.mark.asyncio +async def test_pass_through_request_streamed_response_is_owned_by_the_caller(): + """ + Regression: with passthrough_managed_object_ids on, a streamed + POST /openai_passthrough/v1/responses left the raw resp_ id in the stream and + recorded no owner, so any other key could read, continue, and delete it. + """ + import litellm + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.types.llms.custom_http import httpxSpecialProvider + + raw_id = "resp_0123456789abcdef" + upstream_body = ( + b'event: response.created\ndata: {"type": "response.created", "response": {"id": "%s"}}\n\n' + b'event: response.completed\ndata: {"type": "response.completed", "response": {"id": "%s"}}\n\n' + ) % (raw_id.encode(), raw_id.encode()) + prisma_client = MagicMock() + prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) + prisma_client.db.litellm_managedobjecttable.upsert = AsyncMock(return_value=None) + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) + + def transport_handler(upstream_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, content=upstream_body, headers={"content-type": "text/event-stream"}) + + real_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.PassThroughEndpoint, + params={"timeout": resolve_pass_through_request_timeout(None)}, + ) + cache_dict = litellm.in_memory_llm_clients_cache.cache_dict + cache_key = next(key for key, cached in cache_dict.items() if cached is real_handler) + cache_dict[cache_key] = SimpleNamespace(client=httpx.AsyncClient(transport=httpx.MockTransport(transport_handler))) + + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) + mock_proxy_logging.get_proxy_hook = MagicMock(return_value=MagicMock()) + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.scope = {"path": "/openai_passthrough/v1/responses"} + mock_request.url = MagicMock() + mock_request.url.path = "/openai_passthrough/v1/responses" + mock_request.body = AsyncMock(return_value=b'{"model": "gpt-5.1", "input": "hi", "stream": true}') + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + + flag_on = {"passthrough_managed_object_ids": True} + proxy_server_globals = ( + patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging), # test-quality-ok: read at call time + patch("litellm.proxy.proxy_server.general_settings", flag_on), # test-quality-ok: read at call time + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: read at call time + ) + + try: + with ExitStack() as stack: + for patched_global in proxy_server_globals: + stack.enter_context(patched_global) + response = await pass_through_request( + request=mock_request, + target="https://api.openai.com/v1/responses", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(user_id="user-a", team_id="team-a"), + custom_llm_provider="openai", + ) + streamed = b"".join([chunk async for chunk in response.body_iterator]) + finally: + cache_dict[cache_key] = real_handler + + assert response.status_code == 200 + prisma_client.db.litellm_managedobjecttable.upsert.assert_awaited_once() + created = prisma_client.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]["create"] + assert created["created_by"] == "user-a" + assert created["team_id"] == "team-a" + assert created["model_object_id"] == f"passthrough:openai:{raw_id}" + managed_id = created["unified_object_id"] + assert raw_id.encode() not in streamed + assert streamed == upstream_body.replace(raw_id.encode(), managed_id.encode()) + + @pytest.mark.asyncio async def test_create_pass_through_endpoint(): """ diff --git a/tests/test_litellm/test_responses_id_security.py b/tests/test_litellm/test_responses_id_security.py index 17487030cc1..763ee4dac00 100644 --- a/tests/test_litellm/test_responses_id_security.py +++ b/tests/test_litellm/test_responses_id_security.py @@ -9,8 +9,15 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException -from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.proxy.hooks.responses_id_security import ( + ResponsesIDSecurity, + _is_responses_api_create_route, +) +from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, +) from litellm.types.utils import SpecialEnums @@ -575,6 +582,115 @@ class TestAsyncPreCallHook: assert "team" in exc_info.value.detail.lower() +class TestIsResponsesApiCreateRoute: + """Test the route gate that decides whether a streamed response id is encrypted.""" + + @pytest.mark.parametrize( + "route", + [ + "/v1/responses", + "/responses", + "/openai/v1/responses", + ], + ) + def test_create_routes_match(self, route): + assert _is_responses_api_create_route(route) is True + + @pytest.mark.parametrize( + "route", + [ + None, + "/chat/completions", + "/openai/v1/chat/completions", + "/v1/responses/{response_id}", + "/openai/v1/responses/{response_id}", + "/v1/responsesX", + "/responsesX", + ], + ) + def test_non_create_routes_do_not_match(self, route): + assert _is_responses_api_create_route(route) is False + + +class TestAsyncPostCallStreamingIteratorHook: + """Regression test for LIT-6167: streamed responses on /openai/v1/responses and + /responses must have their ids security-encrypted, not just on the exact + /v1/responses path. A streamed create emits ResponseCompletedEvent, whose + client-visible id lives on event.response.id, so the test drives that production + event shape (not a top-level id) and uses real encryption, asserting the id + round-trips back to the raw provider id plus the caller's user/team, which is the + access-control wrapper the aliases were leaking without.""" + + @staticmethod + async def _agen(chunks): + for chunk in chunks: + yield chunk + + @staticmethod + def _completed_event(response_id): + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id=response_id, + created_at=0, + model="gpt-5.1", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + + async def _drain_streamed_id(self, responses_id_security, route, monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-abcdefghij") + event = self._completed_event("resp_rawprovider123") + + mock_auth = MagicMock() + mock_auth.user_id = "user-a" + mock_auth.team_id = "team-a" + mock_auth.request_route = route + + collected = [ + out + async for out in responses_id_security.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_auth, + response=self._agen([event]), + request_data={}, + ) + ] + return collected[0].response.id + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "route", + ["/v1/responses", "/responses", "/openai/v1/responses"], + ) + async def test_streamed_id_encrypted_on_all_responses_routes( + self, responses_id_security, route, monkeypatch + ): + streamed_id = await self._drain_streamed_id(responses_id_security, route, monkeypatch) + + assert streamed_id != "resp_rawprovider123" + assert responses_id_security._is_encrypted_response_id(streamed_id) + assert responses_id_security._decrypt_response_id(streamed_id) == ( + "resp_rawprovider123", + "user-a", + "team-a", + ) + + @pytest.mark.asyncio + async def test_streamed_id_untouched_on_non_responses_route( + self, responses_id_security, monkeypatch + ): + streamed_id = await self._drain_streamed_id( + responses_id_security, "/chat/completions", monkeypatch + ) + + assert streamed_id == "resp_rawprovider123" + assert not responses_id_security._is_encrypted_response_id(streamed_id) + + class TestAsyncPostCallSuccessHook: """Test async_post_call_success_hook function""" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 5ee6160fbdb..cceb034a20b 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -10607,3 +10607,45 @@ async def test_async_function_with_fallbacks_scrubs_spoofed_values_from_sibling_ assert litellm_metadata["client_key"] == "client_value" assert metadata["attempted_fallbacks"] == 0 assert metadata["original_model_group"] == "gpt-3.5-turbo" + + +def _permission_denied_error() -> litellm.PermissionDeniedError: + return litellm.PermissionDeniedError( + message="OpenrouterException - this key has no access to the model", + llm_provider="openrouter", + model="openrouter/openai/gpt-4o", + response=httpx.Response(status_code=403, request=httpx.Request(method="POST", url="https://openrouter.ai")), + ) + + +def test_permission_denied_error_is_not_retried_against_a_single_deployment(): + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "openrouter/openai/gpt-4o", "api_key": "sk-test"}}, + ] + ) + + with pytest.raises(litellm.PermissionDeniedError): + router.should_retry_this_error( + error=_permission_denied_error(), + healthy_deployments=router.model_list, + all_deployments=router.model_list, + ) + + +def test_permission_denied_error_is_retried_when_other_deployments_exist(): + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "openrouter/openai/gpt-4o", "api_key": "sk-test"}}, + {"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test"}}, + ] + ) + + assert ( + router.should_retry_this_error( + error=_permission_denied_error(), + healthy_deployments=router.model_list, + all_deployments=router.model_list, + ) + is True + )